From d3cd569e940479db7a2aab0a9b12aa81f1206ced Mon Sep 17 00:00:00 2001 From: Brian Clozel Date: Thu, 26 Sep 2024 14:20:44 +0200 Subject: [PATCH] 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 --- .../server/webmvc/GraphQlSseHandler.java | 3 +- .../webmvc/GraphQlWebSocketHandler.java | 1 + .../server/WebSocketHandlerTestSupport.java | 18 +++++++++- .../server/webmvc/GraphQlSseHandlerTests.java | 36 ++++++++++++++++++- .../webmvc/GraphQlWebSocketHandlerTests.java | 24 +++++++++++-- .../server/webmvc/TestWebSocketSession.java | 3 +- 6 files changed, 78 insertions(+), 7 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlSseHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlSseHandler.java index 1f846e26..957dcc7b 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlSseHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlSseHandler.java @@ -106,7 +106,8 @@ public class GraphQlSseHandler extends AbstractGraphQlHttpHandler { this.sseBuilder.data(value); } catch (IOException exception) { - onError(exception); + cancel(); + hookOnError(exception); } } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java index 6dce0114..a87706a3 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java @@ -609,6 +609,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub request(1); } catch (IOException ex) { + cancel(); ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.session, ex, logger); } } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java index d152a8a7..47e738fd 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/WebSocketHandlerTestSupport.java +++ b/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 @@ public abstract class WebSocketHandlerTestSupport { .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"); + } + } + } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlSseHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlSseHandlerTests.java index 49bf9a00..42e810db 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlSseHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlSseHandlerTests.java @@ -21,9 +21,12 @@ import java.io.IOException; 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 org.springframework.web.servlet.function.ServerResponse; 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> 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 @@ class GraphQlSseHandlerTests { """); } + @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)) diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java index 7383ac47..66e3ffd4 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandlerTests.java +++ b/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 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { .verify(TIMEOUT); } + @Test + void brokenPipeShouldCancelPublisher() throws Exception { + this.session = new BrokenPipeSession(); + handle(this.handler, new TextMessage("{\"type\":\"connection_init\"}"), new TextMessage(BOOK_SUBSCRIPTION)); + + BiConsumer, String> bookPayloadAssertion = (message, bookId) -> { + GraphQlWebSocketMessage actual = decode(message); + assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID); + assertThat(actual.resolvedType()).isEqualTo(GraphQlWebSocketMessageType.NEXT); + assertThat(actual.>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 = diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/TestWebSocketSession.java b/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/TestWebSocketSession.java index 59eb9ab1..62ea006a 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/webmvc/TestWebSocketSession.java +++ b/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 class TestWebSocketSession implements WebSocketSession { } @Override - public void sendMessage(WebSocketMessage message) { + public void sendMessage(WebSocketMessage message) throws IOException { emitMessagesSignal(this.messagesSink.tryEmitNext(message)); }