Polishing

See gh-1149
This commit is contained in:
Brian Clozel
2025-03-14 09:38:56 +01:00
parent f3f5516b8a
commit a6c69eb7ac
12 changed files with 202 additions and 136 deletions

View File

@@ -36,12 +36,6 @@ import org.springframework.lang.Nullable;
*/
public interface ExecutionGraphQlRequest extends GraphQlRequest {
/**
* Key of the GraphQL context entry that holds a {@code Mono<Void>} that completes
* when the inbound GraphQL request is cancelled at the transport level.
*/
String CANCEL_PUBLISHER_CONTEXT_KEY = ExecutionGraphQlRequest.class.getName() + ".cancelled";
/**
* Return the transport assigned id for the request that in turn sets
* {@link ExecutionInput.Builder#executionId(ExecutionId) executionId}.

View File

@@ -32,7 +32,7 @@ import reactor.core.publisher.Mono;
import org.springframework.core.CoroutinesUtils;
import org.springframework.core.KotlinDetector;
import org.springframework.data.util.KotlinReflectionUtils;
import org.springframework.graphql.execution.ContextSnapshotFactoryHelper;
import org.springframework.graphql.execution.ContextPropagationHelper;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -153,7 +153,7 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
CompletableFuture<Object> future = new CompletableFuture<>();
this.executor.execute(() -> {
try {
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext);
ContextSnapshot snapshot = ContextPropagationHelper.captureFrom(graphQLContext);
Object value = snapshot.wrap((Callable<?>) result).call();
future.complete(value);
}

View File

@@ -39,7 +39,6 @@ import org.reactivestreams.Publisher;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.graphql.ExecutionGraphQlRequest;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -80,24 +79,22 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
public Object get(DataFetchingEnvironment env) throws Exception {
GraphQLContext graphQlContext = env.getGraphQlContext();
ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.getInstance(graphQlContext);
ContextSnapshotFactory snapshotFactory = ContextPropagationHelper.getInstance(graphQlContext);
ContextSnapshot snapshot = (env.getLocalContext() instanceof GraphQLContext localContext) ?
snapshotFactory.captureFrom(graphQlContext, localContext) :
snapshotFactory.captureFrom(graphQlContext);
Mono<Void> cancelledRequest = graphQlContext.get(ExecutionGraphQlRequest.CANCEL_PUBLISHER_CONTEXT_KEY);
Object value = snapshot.wrap(() -> this.delegate.get(env)).call();
if (value instanceof DataFetcherResult<?> dataFetcherResult) {
Object adapted = updateValue(dataFetcherResult.getData(), snapshot, cancelledRequest);
Object adapted = updateValue(dataFetcherResult.getData(), snapshot, graphQlContext);
value = DataFetcherResult.newResult()
.data(adapted)
.errors(dataFetcherResult.getErrors())
.localContext(dataFetcherResult.getLocalContext()).build();
}
else {
value = updateValue(value, snapshot, cancelledRequest);
value = updateValue(value, snapshot, graphQlContext);
}
return value;
@@ -105,7 +102,7 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
@SuppressWarnings("ReactiveStreamsUnusedPublisher")
private @Nullable Object updateValue(
@Nullable Object value, ContextSnapshot snapshot, @Nullable Mono<Void> cancelledRequest) {
@Nullable Object value, ContextSnapshot snapshot, GraphQLContext graphQlContext) {
if (value == null) {
return null;
@@ -121,19 +118,14 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
return this.subscriptionExceptionResolver.resolveException(exception)
.flatMap((errors) -> Mono.error(new SubscriptionPublisherException(errors, exception)));
});
if (cancelledRequest != null) {
subscriptionResult = subscriptionResult.takeUntilOther(cancelledRequest);
}
return subscriptionResult.contextWrite(snapshot::updateContext);
return ContextPropagationHelper.bindCancelFrom(subscriptionResult, graphQlContext)
.contextWrite(snapshot::updateContext);
}
value = ReactiveAdapterRegistryHelper.toMonoIfReactive(value);
if (value instanceof Mono<?> mono) {
if (cancelledRequest != null) {
mono = mono.takeUntilOther(cancelledRequest);
}
value = mono.contextWrite(snapshot::updateContext).toFuture();
value = ContextPropagationHelper.bindCancelFrom(mono, graphQlContext).contextWrite(snapshot::updateContext).toFuture();
}
return value;

View File

@@ -0,0 +1,170 @@
/*
* Copyright 2020-2025 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 graphql.GraphQLContext;
import io.micrometer.context.ContextSnapshot;
import io.micrometer.context.ContextSnapshotFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
/**
* Helper for propagating context values from and to Reactor and GraphQL contexts.
*
* @author Rossen Stoyanchev
* @author Brian Clozel
* @since 1.3.5
*/
public abstract class ContextPropagationHelper {
private static final ContextSnapshotFactory sharedInstance = ContextSnapshotFactory.builder().build();
private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextPropagationHelper.class.getName() + ".KEY";
private static final String CANCEL_PUBLISHER_KEY = ContextPropagationHelper.class.getName() + ".cancelled";
/**
* Select a {@code ContextSnapshotFactory} instance to use, either the one
* passed in if it is not {@code null}, or a shared, static instance.
* @param factory the candidate factory instance to use if not {@code null}
* @return the instance to use
*/
public static ContextSnapshotFactory selectInstance(@Nullable ContextSnapshotFactory factory) {
if (factory != null) {
return factory;
}
return sharedInstance;
}
/**
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
* @param factory the instance to save
* @param context the context to save the instance to
* @return a new context with the saved instance
*/
public static Context saveInstance(ContextSnapshotFactory factory, Context context) {
return context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
}
/**
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
* @param factory the instance to save
* @param context the context to save the instance to
*/
public static void saveInstance(ContextSnapshotFactory factory, GraphQLContext context) {
context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
}
/**
* Access the {@code ContextSnapshotFactory} from the given {@link ContextView}
* or return a shared, static instance.
* @param contextView the context where the instance is saved
* @return the instance to use
*/
public static ContextSnapshotFactory getInstance(ContextView contextView) {
ContextSnapshotFactory factory = contextView.getOrDefault(CONTEXT_SNAPSHOT_FACTORY_KEY, null);
return selectInstance(factory);
}
/**
* Access the {@code ContextSnapshotFactory} from the given {@link GraphQLContext}
* or return a shared, static instance.
* @param context the context where the instance is saved
* @return the instance to use
*/
public static ContextSnapshotFactory getInstance(GraphQLContext context) {
ContextSnapshotFactory factory = context.get(CONTEXT_SNAPSHOT_FACTORY_KEY);
return selectInstance(factory);
}
/**
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
* capture from the given {@link ContextView}.
* @param contextView the context to capture from
* @return a snapshot from the capture
*/
public static ContextSnapshot captureFrom(ContextView contextView) {
ContextSnapshotFactory factory = getInstance(contextView);
return selectInstance(factory).captureFrom(contextView);
}
/**
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
* capture from the given {@link GraphQLContext}.
* @param context the context to capture from
* @return a snapshot from the capture
*/
public static ContextSnapshot captureFrom(GraphQLContext context) {
ContextSnapshotFactory factory = getInstance(context);
return selectInstance(factory).captureFrom(context);
}
/**
* Create a publisher and store it into the given {@link GraphQLContext}.
* This publisher can then be used to propagate cancel signals to upstream publishers.
* @param context the current GraphQL context
* @since 1.3.5
*/
public static Sinks.Empty<Void> createCancelPublisher(GraphQLContext context) {
Sinks.Empty<Void> requestCancelled = Sinks.empty();
context.put(CANCEL_PUBLISHER_KEY, requestCancelled.asMono());
return requestCancelled;
}
/**
* Bind the source {@link Mono} to the publisher from the given {@link GraphQLContext}.
* The returned {@code Mono} will be cancelled when this publisher completes.
* Subscribers must use the returned {@code Mono} instance.
* @param source the source {@code Mono}
* @param context the current GraphQL context
* @param <T> the type of published elements
* @return the new {@code Mono} that will be cancelled when notified
* @since 1.3.5
*/
public static <T> Mono<T> bindCancelFrom(Mono<T> source, GraphQLContext context) {
Mono<Void> cancelSignal = context.get(CANCEL_PUBLISHER_KEY);
if (cancelSignal != null) {
return source.takeUntilOther(cancelSignal);
}
return source;
}
/**
* Bind the source {@link Flux} to the publisher from the given {@link GraphQLContext}.
* The returned {@code Flux} will be cancelled when this publisher completes.
* Subscribers must use the returned {@code Mono} instance.
* @param source the source {@code Mono}
* @param context the current GraphQL context
* @param <T> the type of published elements
* @return the new {@code Mono} that will be cancelled when notified
* @since 1.3.5
*/
public static <T> Flux<T> bindCancelFrom(Flux<T> source, GraphQLContext context) {
Mono<Void> cancelSignal = context.get(CANCEL_PUBLISHER_KEY);
if (cancelSignal != null) {
return source.takeUntilOther(cancelSignal);
}
return source;
}
}

View File

@@ -16,13 +16,7 @@
package org.springframework.graphql.execution;
import graphql.GraphQLContext;
import io.micrometer.context.ContextSnapshot;
import io.micrometer.context.ContextSnapshotFactory;
import reactor.util.context.Context;
import reactor.util.context.ContextView;
import org.springframework.lang.Nullable;
/**
* Helper to use a single {@link ContextSnapshotFactory} instance by saving and
@@ -30,88 +24,9 @@ import org.springframework.lang.Nullable;
*
* @author Rossen Stoyanchev
* @since 1.3.0
* @deprecated since 1.3.5 in favor of {@link ContextPropagationHelper}.
*/
public abstract class ContextSnapshotFactoryHelper {
private static final ContextSnapshotFactory sharedInstance = ContextSnapshotFactory.builder().build();
private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextSnapshotFactoryHelper.class.getName() + ".KEY";
/**
* Select a {@code ContextSnapshotFactory} instance to use, either the one
* passed in if it is not {@code null}, or a shared, static instance.
* @param factory the candidate factory instance to use if not {@code null}
* @return the instance to use
*/
public static ContextSnapshotFactory selectInstance(@Nullable ContextSnapshotFactory factory) {
if (factory != null) {
return factory;
}
return sharedInstance;
}
/**
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
* @param factory the instance to save
* @param context the context to save the instance to
* @return a new context with the saved instance
*/
public static Context saveInstance(ContextSnapshotFactory factory, Context context) {
return context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
}
/**
* Save the {@code ContextSnapshotFactory} in the given {@link Context}.
* @param factory the instance to save
* @param context the context to save the instance to
*/
public static void saveInstance(ContextSnapshotFactory factory, GraphQLContext context) {
context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory);
}
/**
* Access the {@code ContextSnapshotFactory} from the given {@link ContextView}
* or return a shared, static instance.
* @param contextView the context where the instance is saved
* @return the instance to use
*/
public static ContextSnapshotFactory getInstance(ContextView contextView) {
ContextSnapshotFactory factory = contextView.getOrDefault(CONTEXT_SNAPSHOT_FACTORY_KEY, null);
return selectInstance(factory);
}
/**
* Access the {@code ContextSnapshotFactory} from the given {@link GraphQLContext}
* or return a shared, static instance.
* @param context the context where the instance is saved
* @return the instance to use
*/
public static ContextSnapshotFactory getInstance(GraphQLContext context) {
ContextSnapshotFactory factory = context.get(CONTEXT_SNAPSHOT_FACTORY_KEY);
return selectInstance(factory);
}
/**
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
* capture from the given {@link ContextView}.
* @param contextView the context to capture from
* @return a snapshot from the capture
*/
public static ContextSnapshot captureFrom(ContextView contextView) {
ContextSnapshotFactory factory = getInstance(contextView);
return selectInstance(factory).captureFrom(contextView);
}
/**
* Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to
* capture from the given {@link GraphQLContext}.
* @param context the context to capture from
* @return a snapshot from the capture
*/
public static ContextSnapshot captureFrom(GraphQLContext context) {
ContextSnapshotFactory factory = getInstance(context);
return selectInstance(factory).captureFrom(context);
}
@Deprecated(since = "1.3.5", forRemoval = true)
public abstract class ContextSnapshotFactoryHelper extends ContextPropagationHelper {
}

View File

@@ -96,7 +96,7 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher
private List<GraphQLError> resolveInternal(Throwable exception, DataFetchingEnvironment env) {
try {
return (this.threadLocalContextAware) ?
ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext())
ContextPropagationHelper.captureFrom(env.getGraphQlContext())
.wrap(() -> resolveToMultipleErrors(exception, env))
.call() :
resolveToMultipleErrors(exception, env);

View File

@@ -233,7 +233,7 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
GraphQLContext graphQLContext = environment.getContext();
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext);
ContextSnapshot snapshot = ContextPropagationHelper.captureFrom(graphQLContext);
try {
return snapshot.wrap(() ->
this.loader.apply(keys, environment)
@@ -281,7 +281,7 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
@Override
public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
GraphQLContext graphQLContext = environment.getContext();
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext);
ContextSnapshot snapshot = ContextPropagationHelper.captureFrom(graphQLContext);
try {
return snapshot.wrap(() ->
this.loader.apply(keys, environment)

View File

@@ -97,20 +97,19 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService {
ExecutionInput executionInput = request.toExecutionInput();
ContextSnapshotFactory factory = ContextSnapshotFactoryHelper.getInstance(contextView);
ContextSnapshotFactory factory = ContextPropagationHelper.getInstance(contextView);
GraphQLContext graphQLContext = executionInput.getGraphQLContext();
ContextSnapshotFactoryHelper.saveInstance(factory, graphQLContext);
ContextPropagationHelper.saveInstance(factory, graphQLContext);
factory.captureFrom(contextView).updateContext(graphQLContext);
Sinks.Empty<Void> requestCancelled = Sinks.empty();
graphQLContext.put(ExecutionGraphQlRequest.CANCEL_PUBLISHER_CONTEXT_KEY, requestCancelled.asMono());
ExecutionInput executionInputToUse = registerDataLoaders(executionInput);
Sinks.Empty<Void> cancelPublisher = ContextPropagationHelper.createCancelPublisher(graphQLContext);
return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(executionInputToUse))
.onErrorResume((ex) -> ex instanceof GraphQLError, (ex) ->
Mono.just(ExecutionResult.newExecutionResult().addError((GraphQLError) ex).build()))
.map((result) -> new DefaultExecutionGraphQlResponse(executionInputToUse, result))
.doOnCancel(requestCancelled::tryEmitEmpty);
.doOnCancel(cancelPublisher::tryEmitEmpty);
});
}

View File

@@ -69,7 +69,7 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler
Throwable exception = unwrapException(handlerParameters);
DataFetchingEnvironment env = handlerParameters.getDataFetchingEnvironment();
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext());
ContextSnapshot snapshot = ContextPropagationHelper.captureFrom(env.getGraphQlContext());
try {
return Flux.fromIterable(this.resolvers)
.flatMap((resolver) -> resolver.resolveException(exception, env))

View File

@@ -83,7 +83,7 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti
public final Mono<List<GraphQLError>> resolveException(Throwable exception) {
if (this.threadLocalContextAware) {
return Mono.deferContextual((contextView) -> {
ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(contextView);
ContextSnapshot snapshot = ContextPropagationHelper.captureFrom(contextView);
try {
List<GraphQLError> errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call();
return Mono.justOrEmpty(errors);

View File

@@ -25,7 +25,7 @@ import io.micrometer.context.ContextSnapshotFactory;
import reactor.core.publisher.Mono;
import org.springframework.graphql.ExecutionGraphQlService;
import org.springframework.graphql.execution.ContextSnapshotFactoryHelper;
import org.springframework.graphql.execution.ContextPropagationHelper;
import org.springframework.graphql.server.WebGraphQlInterceptor.Chain;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
@@ -81,7 +81,7 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
@Override
public WebGraphQlHandler build() {
ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.selectInstance(this.snapshotFactory);
ContextSnapshotFactory snapshotFactory = ContextPropagationHelper.selectInstance(this.snapshotFactory);
Chain endOfChain = (request) -> this.service.execute(request).map(WebGraphQlResponse::new);
@@ -107,7 +107,7 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder {
public Mono<WebGraphQlResponse> handleRequest(WebGraphQlRequest request) {
ContextSnapshot snapshot = snapshotFactory.captureAll();
return executionChain.next(request).contextWrite((context) -> {
context = ContextSnapshotFactoryHelper.saveInstance(snapshotFactory, context);
context = ContextPropagationHelper.saveInstance(snapshotFactory, context);
return snapshot.updateContext(context);
});
}

View File

@@ -19,7 +19,6 @@ package org.springframework.graphql.execution;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
@@ -47,7 +46,6 @@ import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.test.StepVerifier;
import org.springframework.graphql.ExecutionGraphQlRequest;
import org.springframework.graphql.GraphQlSetup;
import org.springframework.graphql.ResponseHelper;
import org.springframework.graphql.TestThreadLocalAccessor;
@@ -300,9 +298,8 @@ public class ContextDataFetcherDecoratorTests {
)
.toGraphQl();
Sinks.Empty<Void> requestCancelled = Sinks.empty();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }")
.graphQLContext(Map.of(ExecutionGraphQlRequest.CANCEL_PUBLISHER_CONTEXT_KEY, requestCancelled.asMono())).build();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
requestCancelled.tryEmitEmpty();
@@ -320,9 +317,8 @@ public class ContextDataFetcherDecoratorTests {
)
.toGraphQl();
Sinks.Empty<Void> requestCancelled = Sinks.empty();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }")
.graphQLContext(Map.of(ExecutionGraphQlRequest.CANCEL_PUBLISHER_CONTEXT_KEY, requestCancelled.asMono())).build();
ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
requestCancelled.tryEmitEmpty();
@@ -339,14 +335,14 @@ public class ContextDataFetcherDecoratorTests {
.doOnCancel(() -> dataFetcherCancelled.set(true))
)
.toGraphQl();
Sinks.Empty<Void> requestCancelled = Sinks.empty();
ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }")
.graphQLContext(Map.of(ExecutionGraphQlRequest.CANCEL_PUBLISHER_CONTEXT_KEY, requestCancelled.asMono())).build();
ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }").build();
Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
ExecutionResult executionResult = graphQl.executeAsync(input).get();
ResponseHelper.forSubscription(executionResult).subscribe();
requestCancelled.tryEmitEmpty();
await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
assertThat(dataFetcherCancelled).isTrue();
}