Align cancelling behavior to graphql-java request execution

This commit uses the new `ExecutionInput.cancel()` support from
graphql-java to cancel data fetchers, if not already in flight.

Closes gh-1171

Author: bclozel

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

@@ -22,7 +22,6 @@
 import graphql.ExecutionInput;
 import graphql.GraphQLContext;
 import graphql.TrivialDataFetcher;
-import graphql.execution.AbortExecutionException;
 import graphql.execution.DataFetcherResult;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
@@ -53,7 +52,6 @@
  * <li>Re-establish Reactor Context passed via {@link ExecutionInput}.
  * <li>Re-establish ThreadLocal context passed via {@link ExecutionInput}.
  * <li>Resolve exceptions from a GraphQL subscription {@link Publisher}.
- * <li>Propagate the cancellation signal to {@code DataFetcher} from the transport layer.
  * </ul>
  *
  * @author Rossen Stoyanchev
@@ -109,11 +107,6 @@ private ContextDataFetcherDecorator(
 		if (value == null) {
 			return null;
 		}
-		if (ContextPropagationHelper.isCancelled(graphQlContext)) {
-			return DataFetcherResult.newResult()
-					.error(new AbortExecutionException("GraphQL request has been cancelled by the client."))
-					.build();
-		}
 
 		if (this.subscription) {
 			Flux<?> subscriptionResult = ReactiveAdapterRegistryHelper.toSubscriptionFlux(value)
@@ -125,8 +118,7 @@ private ContextDataFetcherDecorator(
 						return this.subscriptionExceptionResolver.resolveException(exception)
 								.flatMap((errors) -> Mono.error(new SubscriptionPublisherException(errors, exception)));
 					});
-			return ContextPropagationHelper.bindCancelFrom(subscriptionResult, graphQlContext)
-					.contextWrite(snapshot::updateContext);
+			return subscriptionResult.contextWrite(snapshot::updateContext);
 		}
 
 		value = ReactiveAdapterRegistryHelper.toMonoIfReactive(value);

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

@@ -16,15 +16,10 @@
 
 package org.springframework.graphql.execution;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 import graphql.GraphQLContext;
 import io.micrometer.context.ContextSnapshot;
 import io.micrometer.context.ContextSnapshotFactory;
 import org.jspecify.annotations.Nullable;
-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;
 
@@ -41,10 +36,6 @@ public abstract class ContextPropagationHelper {
 
 	private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextPropagationHelper.class.getName() + ".KEY";
 
-	private static final String CANCELED_KEY = ContextPropagationHelper.class.getName() + ".canceled";
-
-	private static final String CANCELED_PUBLISHER_KEY = ContextPropagationHelper.class.getName() + ".canceledPublisher";
-
 
 	/**
 	 * Select a {@code ContextSnapshotFactory} instance to use, either the one
@@ -122,54 +113,4 @@ public static ContextSnapshot captureFrom(GraphQLContext context) {
 		return selectInstance(factory).captureFrom(context);
 	}
 
-	/**
-	 * Create an atomic boolean and store it into the given {@link GraphQLContext}.
-	 * This boolean value can then be checked by upstream publishers to know whether the request is canceled.
-	 * @param context the current GraphQL context
-	 * @since 1.3.6
-	 */
-	public static Runnable createCancelSignal(GraphQLContext context) {
-		AtomicBoolean requestCancelled = new AtomicBoolean();
-		Sinks.Empty<Void> cancelSignal = Sinks.empty();
-		context.put(CANCELED_KEY, requestCancelled);
-		context.put(CANCELED_PUBLISHER_KEY, cancelSignal.asMono());
-		return () -> {
-			requestCancelled.set(true);
-			cancelSignal.tryEmitEmpty();
-		};
-	}
-
-	/**
-	 * Return {@code true} if the current request has been cancelled, {@code false} otherwise.
-	 * This checks whether a {@link #createCancelSignal(GraphQLContext) cancellation publisher is present}
-	 * in the given context and the cancel signal has fired already.
-	 * @param context the current GraphQL context
-	 * @since 1.4.0
-	 */
-	public static boolean isCancelled(GraphQLContext context) {
-		AtomicBoolean requestCancelled = context.get(CANCELED_KEY);
-		if (requestCancelled != null) {
-			return requestCancelled.get();
-		}
-		return false;
-	}
-
-	/**
-	 * 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(CANCELED_PUBLISHER_KEY);
-		if (cancelSignal != null) {
-			return source.takeUntilOther(cancelSignal);
-		}
-		return source;
-	}
-
 }

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

@@ -91,13 +91,12 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 			factory.captureFrom(contextView).updateContext(graphQLContext);
 
 			ExecutionInput executionInputToUse = registerDataLoaders(executionInput);
-			Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(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(cancelSignal::run);
+					.doOnCancel(executionInputToUse::cancel);
 		});
 	}
 

spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java

@@ -42,7 +42,6 @@
 import io.micrometer.context.ContextRegistry;
 import io.micrometer.context.ContextSnapshot;
 import io.micrometer.context.ContextSnapshotFactory;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -291,60 +290,19 @@ void trivialDataFetcherIsNotDecorated() {
 		assertThat(dataFetcher).isInstanceOf(TrivialDataFetcher.class);
 	}
 
-	@Test
-	@Disabled("until https://github.com/spring-projects/spring-graphql/issues/1171")
-	void cancelMonoDataFetcherWhenRequestCancelled() {
-		AtomicBoolean dataFetcherCancelled = new AtomicBoolean();
-		GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
-				.queryFetcher("greeting", (env) ->
-						Mono.just("Hello")
-								.delayElement(Duration.ofSeconds(1))
-								.doOnCancel(() -> dataFetcherCancelled.set(true))
-						)
-				.toGraphQl();
-
-		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
-
-		CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
-		cancelSignal.run();
-		await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
-	}
-
-	@Test
-	@Disabled("until https://github.com/spring-projects/spring-graphql/issues/1171")
-	void cancelFluxDataFetcherWhenRequestCancelled() {
-		AtomicBoolean dataFetcherCancelled = new AtomicBoolean();
-		GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
-				.queryFetcher("greeting", (env) ->
-						Flux.just("Hello")
-								.delayElements(Duration.ofSeconds(1))
-								.doOnCancel(() -> dataFetcherCancelled.set(true))
-				)
-				.toGraphQl();
-
-		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
-
-		CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
-		cancelSignal.run();
-		await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
-	}
-
 	@Test
 	void returnAbortExecutionForBlockingDataFetcherWhenRequestCancelled() throws Exception {
 		GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
 				.queryFetcher("greeting", (env) -> "Hello")
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
-		cancelSignal.run();
+		input.cancel();
 		ExecutionResult result = graphQl.executeAsync(input).get();
 
 		assertThat(result.getErrors()).hasSize(1);
 		assertThat(result.getErrors().get(0)).isInstanceOf(AbortExecutionException.class)
-				.extracting("message").asString().isEqualTo("GraphQL request has been cancelled by the client.");
+				.extracting("message").asString().isEqualTo("Execution has been asked to be cancelled");
 	}
 
 	@Test
@@ -359,12 +317,10 @@ void cancelFluxDataFetcherSubscriptionWhenRequestCancelled() throws Exception {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }").build();
-		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
-
 		ExecutionResult executionResult = graphQl.executeAsync(input).get();
-		ResponseHelper.forSubscription(executionResult).subscribe();
-		cancelSignal.run();
-
+		input.cancel();
+		StepVerifier.create(ResponseHelper.forSubscription(executionResult))
+				.verifyError(AbortExecutionException.class);
 		await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
 		assertThat(dataFetcherCancelled).isTrue();
 	}

spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultExecutionGraphQlServiceTests.java

@@ -22,14 +22,14 @@
 
 import graphql.ErrorType;
 import org.dataloader.DataLoaderRegistry;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
+import org.springframework.graphql.BookSource;
 import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
 import org.springframework.graphql.GraphQlSetup;
@@ -83,20 +83,24 @@ void shouldHandleGraphQlErrors() {
 	}
 
 	@Test
-	@Disabled("until https://github.com/spring-projects/spring-graphql/issues/1171")
-	void cancellationSupport() {
-		AtomicBoolean cancelled = new AtomicBoolean();
-		Mono<String> greetingMono = Mono.just("hi")
-				.delayElement(Duration.ofSeconds(3))
-				.doOnCancel(() -> cancelled.set(true));
-
-		Mono<ExecutionGraphQlResponse> execution = GraphQlSetup.schemaContent("type Query { greeting: String }")
-				.queryFetcher("greeting", (env) -> greetingMono)
+	void cancellationSupport() throws Exception {
+		AtomicBoolean called = new AtomicBoolean();
+		Mono<Book> bookMono = Mono.just(BookSource.getBookWithoutAuthor(1L))
+				.delayElement(Duration.ofSeconds(1));
+
+		Mono<ExecutionGraphQlResponse> execution = GraphQlSetup.schemaResource(BookSource.schema)
+				.queryFetcher("bookById", (env) -> bookMono)
+				.dataFetcher("Book", "author", (env) -> {
+					called.set(true);
+					return BookSource.getAuthor(1L);
+				})
 				.toGraphQlService()
-				.execute(TestExecutionRequest.forDocument("{ greeting }"));
 
-		StepVerifier.create(execution).thenCancel().verify();
-		assertThat(cancelled).isTrue();
+				.execute(TestExecutionRequest.forDocument("{ bookById(id: 1) { author { firstName } } }"));
+
+		StepVerifier.create(execution).thenAwait(Duration.ofMillis(500)).thenCancel().verify();
+		Thread.sleep(1000);
+		assertThat(called).isFalse();
 	}
 
 }