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 7bff6ee8..12206221 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 @@ -345,12 +345,6 @@ private GraphQlWebSocketMessage decode(TextMessage message) throws IOException { .read(GraphQlWebSocketMessage.class, new HttpInputMessageAdapter(message)); } - private SessionState getSessionInfo(WebSocketSession session) { - SessionState info = this.sessionInfoMap.get(session.getId()); - Assert.notNull(info, "No SessionInfo for " + session); - return info; - } - @SuppressWarnings("unchecked") private Flux handleResponse(WebSocketSession session, String id, WebGraphQlResponse response) { if (logger.isDebugEnabled()) { @@ -364,10 +358,15 @@ private Flux handleResponse(WebSocketSession session, String id, We responseFlux = Flux.from((Publisher) response.getData()) .map(ExecutionResult::toSpecification) .doOnSubscribe((subscription) -> { - Subscription prev = getSessionInfo(session).getSubscriptions().putIfAbsent(id, subscription); - if (prev != null) { - throw new SubscriptionExistsException(); - } + SessionState state = this.sessionInfoMap.get(session.getId()); + if (state == null) { + subscription.cancel(); + return; + } + Subscription prev = state.getSubscriptions().putIfAbsent(id, subscription); + if (prev != null) { + throw new SubscriptionExistsException(); + } }); } else { 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 ec4bd254..204472ae 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 @@ -38,6 +38,7 @@ 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.aot.hint.RuntimeHints; @@ -290,6 +291,25 @@ void messageAfterConnectionClosed() throws Exception { this.handler.handleTextMessage(this.session, new TextMessage("{\"type\":\"ping\"}"))); } + @Test // gh-1501 + void responseAfterConnectionClosed() throws Exception { + Sinks.Empty responseDelay = Sinks.empty(); + GraphQlWebSocketHandler handler = initWebSocketHandler( + (request, chain) -> chain.next(request).delayUntil((response) -> responseDelay.asMono())); + + handle(handler, new TextMessage("{\"type\":\"connection_init\"}"), new TextMessage(BOOK_SUBSCRIPTION)); + handler.afterConnectionClosed(this.session, CloseStatus.NORMAL); + + assertThatNoException().isThrownBy(responseDelay::tryEmitEmpty); + + assertThat(this.session.isOpen()).isTrue(); + StepVerifier.create(this.session.getOutput()) + .consumeNextWith((message) -> assertMessageType(message, GraphQlWebSocketMessageType.CONNECTION_ACK)) + .then(this.session::close) // Complete output Flux + .expectComplete() + .verify(TIMEOUT); + } + @Test void connectionInitRejected() throws Exception {