Merge branch '1.4.x'

Author: bclozel

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

@@ -131,7 +131,7 @@ private ContextDataFetcherDecorator(
 		value = ReactiveAdapterRegistryHelper.toMonoIfReactive(value);
 
 		if (value instanceof Mono<?> mono) {
-			value = ContextPropagationHelper.bindCancelFrom(mono, graphQlContext).contextWrite(snapshot::updateContext).toFuture();
+			value = mono.contextWrite(snapshot::updateContext).toFuture();
 		}
 
 		return value;

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

@@ -16,6 +16,8 @@
 
 package org.springframework.graphql.execution;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import graphql.GraphQLContext;
 import io.micrometer.context.ContextSnapshot;
 import io.micrometer.context.ContextSnapshotFactory;
@@ -39,7 +41,9 @@ public abstract class ContextPropagationHelper {
 
 	private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextPropagationHelper.class.getName() + ".KEY";
 
-	private static final String CANCEL_PUBLISHER_KEY = ContextPropagationHelper.class.getName() + ".cancelled";
+	private static final String CANCELED_KEY = ContextPropagationHelper.class.getName() + ".canceled";
+
+	private static final String CANCELED_PUBLISHER_KEY = ContextPropagationHelper.class.getName() + ".canceledPublisher";
 
 
 	/**
@@ -119,50 +123,37 @@ public static ContextSnapshot captureFrom(GraphQLContext 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.
+	 * 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.5
+	 * @since 1.3.6
 	 */
-	public static Sinks.Empty<Void> createCancelPublisher(GraphQLContext context) {
-		Sinks.Empty<Void> requestCancelled = Sinks.empty();
-		context.put(CANCEL_PUBLISHER_KEY, requestCancelled.asMono());
-		return requestCancelled;
+	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 #createCancelPublisher(GraphQLContext) cancellation publisher is present}
+	 * 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) {
-		Mono<Void> cancelSignal = context.get(CANCEL_PUBLISHER_KEY);
-		if (cancelSignal != null) {
-			return cancelSignal.toFuture().isDone();
+		AtomicBoolean requestCancelled = context.get(CANCELED_KEY);
+		if (requestCancelled != null) {
+			return requestCancelled.get();
 		}
 		return false;
 	}
 
-	/**
-	 * 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.
@@ -174,7 +165,7 @@ public static <T> Mono<T> bindCancelFrom(Mono<T> source, GraphQLContext context)
 	 * @since 1.3.5
 	 */
 	public static <T> Flux<T> bindCancelFrom(Flux<T> source, GraphQLContext context) {
-		Mono<Void> cancelSignal = context.get(CANCEL_PUBLISHER_KEY);
+		Mono<Void> cancelSignal = context.get(CANCELED_PUBLISHER_KEY);
 		if (cancelSignal != null) {
 			return source.takeUntilOther(cancelSignal);
 		}

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

@@ -31,7 +31,6 @@
 import org.dataloader.DataLoaderRegistry;
 import org.jspecify.annotations.Nullable;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Sinks;
 
 import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
@@ -92,13 +91,13 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 			factory.captureFrom(contextView).updateContext(graphQLContext);
 
 			ExecutionInput executionInputToUse = registerDataLoaders(executionInput);
-			Sinks.Empty<Void> cancelPublisher = ContextPropagationHelper.createCancelPublisher(graphQLContext);
+			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(cancelPublisher::tryEmitEmpty);
+					.doOnCancel(cancelSignal::run);
 		});
 	}
 

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

@@ -42,10 +42,10 @@
 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;
-import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
 
 import org.springframework.graphql.GraphQlSetup;
@@ -292,6 +292,7 @@ void trivialDataFetcherIsNotDecorated() {
 	}
 
 	@Test
+	@Disabled("until https://github.com/spring-projects/spring-graphql/issues/1171")
 	void cancelMonoDataFetcherWhenRequestCancelled() {
 		AtomicBoolean dataFetcherCancelled = new AtomicBoolean();
 		GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
@@ -303,14 +304,15 @@ void cancelMonoDataFetcherWhenRequestCancelled() {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
+		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
 
 		CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
-		requestCancelled.tryEmitEmpty();
+		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)
@@ -322,10 +324,10 @@ void cancelFluxDataFetcherWhenRequestCancelled() {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
+		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
 
 		CompletableFuture<ExecutionResult> asyncResult = graphQl.executeAsync(input);
-		requestCancelled.tryEmitEmpty();
+		cancelSignal.run();
 		await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
 	}
 
@@ -336,8 +338,8 @@ void returnAbortExecutionForBlockingDataFetcherWhenRequestCancelled() throws Exc
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
-		requestCancelled.tryEmitEmpty();
+		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
+		cancelSignal.run();
 		ExecutionResult result = graphQl.executeAsync(input).get();
 
 		assertThat(result.getErrors()).hasSize(1);
@@ -357,11 +359,11 @@ void cancelFluxDataFetcherSubscriptionWhenRequestCancelled() throws Exception {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }").build();
-		Sinks.Empty<Void> requestCancelled = ContextPropagationHelper.createCancelPublisher(input.getGraphQLContext());
+		Runnable cancelSignal = ContextPropagationHelper.createCancelSignal(input.getGraphQLContext());
 
 		ExecutionResult executionResult = graphQl.executeAsync(input).get();
 		ResponseHelper.forSubscription(executionResult).subscribe();
-		requestCancelled.tryEmitEmpty();
+		cancelSignal.run();
 
 		await().atMost(Duration.ofSeconds(2)).until(dataFetcherCancelled::get);
 		assertThat(dataFetcherCancelled).isTrue();

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

@@ -22,6 +22,7 @@
 
 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;
@@ -82,6 +83,7 @@ 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")