SSE and WS handlers cancel when client disconnects

Prior to this commit, Server Sent Events and WebSocket handlers for MVC
would not behave properly in case the client disconnects while the
response is written to. This would throw an `IOException` processed by
the `BaseSubscriber` handlers, but would not actively cancel the
upstream publisher. This means the publisher would keep publishing
values even though it is not possible to write to the connection
anymore.

This commit ensures that any exception triggers a cancel signal sent to
the upstream publisher to avoid such cases.

Fixes gh-1060

Author: bclozel

spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlSseHandler.java

@@ -106,7 +106,8 @@ private void writeResult(Map<String, Object> value) {
 				this.sseBuilder.data(value);
 			}
 			catch (IOException exception) {
-				onError(exception);
+				cancel();
+				hookOnError(exception);
 			}
 		}
 

spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java

@@ -609,6 +609,7 @@ protected void hookOnNext(TextMessage nextMessage) {
 				request(1);
 			}
 			catch (IOException ex) {
+				cancel();
 				ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.session, ex, logger);
 			}
 		}

spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java

@@ -16,10 +16,15 @@
 
 package org.springframework.graphql.server;
 
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import reactor.core.publisher.Flux;
 
 import org.springframework.graphql.BookSource;
 import org.springframework.graphql.GraphQlSetup;
+import org.springframework.graphql.server.webmvc.TestWebSocketSession;
+import org.springframework.web.socket.WebSocketMessage;
 
 public abstract class WebSocketHandlerTestSupport {
 
@@ -29,6 +34,8 @@ public abstract class WebSocketHandlerTestSupport {
 
 	protected static final String BOOK_QUERY_PAYLOAD;
 
+	protected static AtomicBoolean SUBSCRIPTION_CANCELLED = new AtomicBoolean();
+
 	static {
 		BOOK_QUERY_PAYLOAD = "{\"query\": \"" +
 				"  query TestQuery {" +
@@ -73,10 +80,19 @@ protected WebGraphQlHandler initHandler(WebGraphQlInterceptor... interceptors) {
 				.subscriptionFetcher("bookSearch", environment -> {
 					String author = environment.getArgument("author");
 					return Flux.fromIterable(BookSource.books())
-							.filter((book) -> book.getAuthor().getFullName().contains(author));
+							.filter((book) -> book.getAuthor().getFullName().contains(author))
+							.doOnCancel(() -> SUBSCRIPTION_CANCELLED.set(true));
 				})
 				.interceptor(interceptors)
 				.toWebGraphQlHandler();
 	}
 
+	public class BrokenPipeSession extends TestWebSocketSession {
+
+		@Override
+		public void sendMessage(WebSocketMessage<?> message) throws IOException {
+			throw new IOException("broken pipe");
+		}
+	}
+
 }

spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlSseHandlerTests.java

@@ -21,9 +21,12 @@
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import graphql.schema.DataFetcher;
 import jakarta.servlet.ServletException;
+import jakarta.servlet.ServletOutputStream;
+import jakarta.servlet.http.HttpServletResponse;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 
@@ -40,6 +43,10 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.awaitility.Awaitility.await;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.BDDMockito.willThrow;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link GraphQlSseHandler}.
@@ -51,9 +58,13 @@ class GraphQlSseHandlerTests {
 	private static final List<HttpMessageConverter<?>> MESSAGE_READERS =
 			List.of(new MappingJackson2HttpMessageConverter());
 
+	private static final AtomicBoolean DATA_FETCHER_CANCELLED = new AtomicBoolean();
+
 	private static final DataFetcher<?> SEARCH_DATA_FETCHER = env -> {
 		String author = env.getArgument("author");
-		return Flux.fromIterable(BookSource.books()).filter((book) -> book.getAuthor().getFullName().contains(author));
+		return Flux.fromIterable(BookSource.books())
+				.filter((book) -> book.getAuthor().getFullName().contains(author))
+				.doOnCancel(() -> DATA_FETCHER_CANCELLED.set(true));
 	};
 
 
@@ -122,6 +133,29 @@ void shouldWriteEventsAndTerminalError() throws Exception {
 				""");
 	}
 
+	@Test
+	void shouldCancelDataFetcherPublisherWhenWritingFails() throws Exception {
+		GraphQlSseHandler handler = createSseHandler(SEARCH_DATA_FETCHER);
+		MockHttpServletRequest servletRequest = createServletRequest("""
+				{ "query": "subscription TestSubscription { bookSearch(author:\\\"Orwell\\\") { id name } }" }
+				""");
+		HttpServletResponse servletResponse = mock(HttpServletResponse.class);
+		ServletOutputStream outputStream = mock(ServletOutputStream.class);
+
+		willThrow(new IOException("broken pipe")).given(outputStream).write(any());
+		given(servletResponse.getOutputStream()).willReturn(outputStream);
+
+		ServerRequest request = ServerRequest.create(servletRequest, MESSAGE_READERS);
+		ServerResponse response = handler.handleRequest(request);
+		if (response instanceof AsyncServerResponse asyncResponse) {
+			asyncResponse.block();
+		}
+
+		response.writeTo(servletRequest, servletResponse, new DefaultContext());
+		await().atMost(Duration.ofMillis(500)).until(DATA_FETCHER_CANCELLED::get);
+
+	}
+
 	private GraphQlSseHandler createSseHandler(DataFetcher<?> dataFetcher) {
 		return new GraphQlSseHandler(GraphQlSetup.schemaResource(BookSource.schema)
 				.queryFetcher("bookById", (env) -> BookSource.getBookWithoutAuthor(1L))

spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java

@@ -77,11 +77,10 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 	private static final Duration TIMEOUT = Duration.ofSeconds(5);
 
-
-	private final TestWebSocketSession session = new TestWebSocketSession();
-
 	private final GraphQlWebSocketHandler handler = initWebSocketHandler();
 
+	private TestWebSocketSession session = new TestWebSocketSession();
+
 
 	@Test
 	void query() throws Exception {
@@ -130,6 +129,25 @@ void subscription() throws Exception {
 				.verify(TIMEOUT);
 	}
 
+	@Test
+	void brokenPipeShouldCancelPublisher() throws Exception {
+		this.session = new BrokenPipeSession();
+		handle(this.handler, new TextMessage("{\"type\":\"connection_init\"}"), new TextMessage(BOOK_SUBSCRIPTION));
+
+		BiConsumer<WebSocketMessage<?>, String> bookPayloadAssertion = (message, bookId) -> {
+			GraphQlWebSocketMessage actual = decode(message);
+			assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+			assertThat(actual.resolvedType()).isEqualTo(GraphQlWebSocketMessageType.NEXT);
+			assertThat(actual.<Map<String, Object>>getPayload())
+					.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.extractingByKey("bookSearch", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.containsEntry("id", bookId);
+		};
+
+		StepVerifier.create(session.getOutput()).verifyComplete();
+		assertThat(SUBSCRIPTION_CANCELLED).isTrue();
+	}
+
 	@Test
 	void keepAlive() throws Exception {
 		GraphQlWebSocketHandler webSocketHandler =

spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/TestWebSocketSession.java

@@ -16,6 +16,7 @@
 
 package org.springframework.graphql.server.webmvc;
 
+import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.net.URI;
 import java.security.Principal;
@@ -119,7 +120,7 @@ public List<WebSocketExtension> getExtensions() {
 	}
 
 	@Override
-	public void sendMessage(WebSocketMessage<?> message) {
+	public void sendMessage(WebSocketMessage<?> message) throws IOException {
 		emitMessagesSignal(this.messagesSink.tryEmitNext(message));
 	}