From d2eaf3c2bcd598cf962d9691c8e84f124e60d0d0 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Fri, 27 Oct 2023 16:40:41 +0100 Subject: [PATCH] Fix reconnect issue in WebSocketGraphQlTransport Closes gh-826 --- .../client/WebSocketGraphQlTransport.java | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java b/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java index f6e1067a..9c5a93c2 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java @@ -96,10 +96,15 @@ final class WebSocketGraphQlTransport implements GraphQlTransport { return Mono.error(new IllegalStateException("WebSocketGraphQlTransport has been stopped")); } - client.execute(uri, headers, handler) - .subscribe(aVoid -> {}, handler::handleWebSocketSessionError, () -> {}); + // Get the session Mono before connecting + Mono sessionMono = handler.getGraphQlSession(); - return handler.getGraphQlSession(); + client.execute(uri, headers, handler) + .subscribe(aVoid -> {}, + handler::handleWebSocketSessionError, + handler::handleWebSocketSessionClosed); + + return sessionMono; }); } @@ -311,10 +316,6 @@ final class WebSocketGraphQlTransport implements GraphQlTransport { } graphQlSession.terminateRequests(closeStatusMessage, closeStatus); }) - .doOnTerminate(() -> { - // Reset GraphQlSession sink to be ready to connect again - this.graphQlSessionSink = Sinks.unsafe().one(); - }) .subscribe(); } @@ -356,6 +357,15 @@ final class WebSocketGraphQlTransport implements GraphQlTransport { } this.graphQlSessionSink.tryEmitError(ex); + this.graphQlSessionSink = Sinks.unsafe().one(); + } + + /** + * This must be called from code that calls the {@code WebSocketClient} + * when execution completes. + */ + public void handleWebSocketSessionClosed() { + this.graphQlSessionSink = Sinks.unsafe().one(); } }