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)); }