Fix high memory usage in ContextPropagationHelper

Prior to this commit, gh-1149 added support for cancellation detection
at the transport level and the propagation of the CANCEL signal to data
fetcher `Publisher`. `ContextPropagationHelper` stores in the context a
`Sink` that emits the cancel signal. All upstream publishers are then
subscribing to this and canceling themselves if they receive a signal.

In the case of high field count queries + high request/sec services,
this would create a significant memory overhead.

This commit reduces the support to the essential parts:
* avoid further calls to data fetchers if the request is canceled
* cancel subscription publishers

This means other in-flight data fetchers that are publisher-based are
not canceled anymore. This should be a good compromise for high RPS

See gh-1242

Author: bclozel

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

@@ -131,7 +131,7 @@ public Object get(DataFetchingEnvironment env) throws Exception {
 		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;
@@ -40,7 +42,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";
 
 
 	/**
@@ -120,50 +124,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.
@@ -175,7 +166,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

@@ -30,7 +30,6 @@
 import io.micrometer.context.ContextSnapshotFactory;
 import org.dataloader.DataLoaderRegistry;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Sinks;
 
 import org.springframework.graphql.ExecutionGraphQlRequest;
 import org.springframework.graphql.ExecutionGraphQlResponse;
@@ -93,13 +92,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")