From 736c33aa4e25133f4e7f74a9abd101a5de21ce99 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Tue, 22 Feb 2022 22:19:37 +0000 Subject: [PATCH] Add GraphQlClient support for WebSocket See gh-10 --- .../client/SubscriptionErrorException.java | 44 ++ .../client/WebSocketCodecDelegate.java | 94 +++ .../client/WebSocketGraphQlTransport.java | 746 ++++++++++++++++++ .../MockWebSocketGraphQlTransportTests.java | 332 ++++++++ .../graphql/client/MockWebSocketServer.java | 200 +++++ 5 files changed, 1416 insertions(+) create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/client/SubscriptionErrorException.java create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketCodecDelegate.java create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java create mode 100644 spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketServer.java diff --git a/spring-graphql/src/main/java/org/springframework/graphql/client/SubscriptionErrorException.java b/spring-graphql/src/main/java/org/springframework/graphql/client/SubscriptionErrorException.java new file mode 100644 index 00000000..1b12df31 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/client/SubscriptionErrorException.java @@ -0,0 +1,44 @@ +/* + * Copyright 2002-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.graphql.client; + +import java.util.List; + +import graphql.GraphQLError; + +/** + * + * @author Rossen Stoyanchev + * @since 1.0.0 + */ +@SuppressWarnings("serial") +public class SubscriptionErrorException extends RuntimeException { + + private final List errors; + + + public SubscriptionErrorException(List errors) { + super("GraphQL subscription error: " + errors); + this.errors = errors; + } + + + public List getErrors() { + return this.errors; + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketCodecDelegate.java b/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketCodecDelegate.java new file mode 100644 index 00000000..cd8cf79a --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketCodecDelegate.java @@ -0,0 +1,94 @@ +/* + * Copyright 2002-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.graphql.client; + +import org.springframework.core.ResolvableType; +import org.springframework.core.codec.Decoder; +import org.springframework.core.codec.Encoder; +import org.springframework.core.io.buffer.DataBuffer; +import org.springframework.core.io.buffer.DataBufferUtils; +import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage; +import org.springframework.http.MediaType; +import org.springframework.http.codec.ClientCodecConfigurer; +import org.springframework.http.codec.CodecConfigurer; +import org.springframework.http.codec.DecoderHttpMessageReader; +import org.springframework.http.codec.EncoderHttpMessageWriter; +import org.springframework.util.Assert; +import org.springframework.util.MimeTypeUtils; +import org.springframework.web.reactive.socket.WebSocketMessage; +import org.springframework.web.reactive.socket.WebSocketSession; + +/** + * Delegate that can be embedded in a class to help with encoding and decoding + * GraphQL over WebSocket messages. + * + * @author Rossen Stoyanchev + * @since 1.0.0 + */ +final class WebSocketCodecDelegate { + + private static final ResolvableType MESSAGE_TYPE = ResolvableType.forClass(GraphQlWebSocketMessage.class); + + + private final Decoder decoder; + + private final Encoder encoder; + + + WebSocketCodecDelegate() { + this(ClientCodecConfigurer.create()); + } + + WebSocketCodecDelegate(CodecConfigurer codecConfigurer) { + Assert.notNull(codecConfigurer, "CodecConfigurer is required"); + this.decoder = initDecoder(codecConfigurer); + this.encoder = initEncoder(codecConfigurer); + } + + private static Decoder initDecoder(CodecConfigurer configurer) { + return configurer.getReaders().stream() + .filter((reader) -> reader.canRead(MESSAGE_TYPE, MediaType.APPLICATION_JSON)) + .map((reader) -> ((DecoderHttpMessageReader) reader).getDecoder()) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("No JSON Decoder")); + } + + private static Encoder initEncoder(CodecConfigurer configurer) { + return configurer.getWriters().stream() + .filter((writer) -> writer.canWrite(MESSAGE_TYPE, MediaType.APPLICATION_JSON)) + .map((writer) -> ((EncoderHttpMessageWriter) writer).getEncoder()) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("No JSON Encoder")); + } + + + @SuppressWarnings("unchecked") + public WebSocketMessage encode(WebSocketSession session, GraphQlWebSocketMessage message) { + + DataBuffer buffer = ((Encoder) this.encoder).encodeValue( + (T) message, session.bufferFactory(), MESSAGE_TYPE, MimeTypeUtils.APPLICATION_JSON, null); + + return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer); + } + + @SuppressWarnings("ConstantConditions") + public GraphQlWebSocketMessage decode(WebSocketMessage webSocketMessage) { + DataBuffer buffer = DataBufferUtils.retain(webSocketMessage.getPayload()); + return (GraphQlWebSocketMessage) this.decoder.decode(buffer, MESSAGE_TYPE, null, null); + } + + +} 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 new file mode 100644 index 00000000..bbcc82ec --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/client/WebSocketGraphQlTransport.java @@ -0,0 +1,746 @@ +/* + * Copyright 2002-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.graphql.client; + +import java.net.URI; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import graphql.ExecutionResult; +import graphql.GraphQLError; +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import reactor.core.Scannable; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; + +import org.springframework.graphql.RequestInput; +import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage; +import org.springframework.http.HttpHeaders; +import org.springframework.http.codec.ClientCodecConfigurer; +import org.springframework.http.codec.CodecConfigurer; +import org.springframework.lang.Nullable; +import org.springframework.util.Assert; +import org.springframework.web.reactive.socket.CloseStatus; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.WebSocketSession; +import org.springframework.web.reactive.socket.client.WebSocketClient; + +/** + * {@link GraphQlTransport} for GraphQL over WebSocket via {@link WebSocketClient}. + * + *

Use the builder to initialize the transport and the {@link GraphQlClient} + * in a single chain: + * + * 

+ * GraphQlClient client =
+ * 		WebSocketGraphQlTransport.builder(url, webSocketClient)
+ * 				.headers(headers -> ... )
+ * 				.buildClient();
+ *
+ * + *

Or build the transport and the client separately: + * + * 

+ * WebSocketGraphQlTransport transport =
+ * 		WebSocketGraphQlTransport.builder(url, webSocketClient)
+ * 				.headers(headers -> ... )
+ * 				.build();
+ *
+ * GraphQlClient client = GraphQlClient.create(transport);
+ *
+ * + *

Once the client is built, you can obtain the underlying transport, mutate + * it, and rebuild transport and client as follows: + * + * 

+ * WebSocketGraphQlTransport transport = client.getTransport(WebSocketGraphQlTransport.class);
+ *
+ * if (transport != null) {
+ * 	GraphQlClient newClient = transport.mutate()
+ *			.headers(headers -> ... )
+ * 			.buildClient();
+ * }
+ *
+ * + * + * @author Rossen Stoyanchev + * @since 1.0.0 + * @see GraphQL over WebSocket protocol + */ +public final class WebSocketGraphQlTransport implements GraphQlTransport { + + private static final Log logger = LogFactory.getLog(WebSocketGraphQlTransport.class); + + + private final GraphQlSessionHandler graphQlSessionHandler; + + private final Mono graphQlSessionMono; + + private final Supplier mutateBuilder; + + + private WebSocketGraphQlTransport( + URI uri, HttpHeaders headers, WebSocketClient client, CodecConfigurer codecConfigurer, + @Nullable Object connectionInitPayload, Consumer> connectionAckHandler, + Supplier mutateBuilder) { + + this.graphQlSessionHandler = new GraphQlSessionHandler( + codecConfigurer, connectionInitPayload, connectionAckHandler); + + this.graphQlSessionMono = initGraphQlSession(uri, headers, client, this.graphQlSessionHandler) + .cacheInvalidateWhen(GraphQlSession::notifyWhenClosed); + + this.mutateBuilder = mutateBuilder; + } + + private static Mono initGraphQlSession( + URI uri, HttpHeaders headers, WebSocketClient client, GraphQlSessionHandler handler) { + + return Mono.defer(() -> { + if (handler.isStopped()) { + return Mono.error(new IllegalStateException("WebSocketGraphQlTransport has been stopped")); + } + + client.execute(uri, headers, handler) + .subscribe(aVoid -> {}, handler::handleWebSocketSessionError, () -> {}); + + return handler.getGraphQlSession(); + }); + } + + + /** + * Start the transport by connecting the WebSocket, sending the + * "connection_init" and waiting for the "connection_ack" message. + * @return {@code Mono} that completes when the WebSocket is connected and + * ready to begin sending GraphQL requests + */ + public Mono start() { + this.graphQlSessionHandler.setStopped(false); + return this.graphQlSessionMono.then(); + } + + /** + * Stop the transport by closing the WebSocket with + * {@link org.springframework.web.reactive.socket.CloseStatus#NORMAL} and + * terminating in-progress requests with an error signal. + *

New requests are rejected from the time of this call. If necessary, + * call {@link #start()} to allow requests again. + * @return {@code Mono} that completes when the underlying session is closed + */ + public Mono stop() { + this.graphQlSessionHandler.setStopped(true); + return this.graphQlSessionMono.flatMap(GraphQlSession::close).onErrorResume(ex -> Mono.empty()); + } + + @Override + public Mono execute(RequestInput input) { + return this.graphQlSessionMono.flatMap(session -> session.execute(input)); + } + + @Override + public Flux executeSubscription(RequestInput input) { + return this.graphQlSessionMono.flatMapMany(session -> session.executeSubscription(input)); + } + + /** + * Create a builder initialized from the configuration of "this" transport. + * Use this to build a new instance configured differently. + */ + public Builder mutate() { + return this.mutateBuilder.get(); + } + + + /** + * Static factory method with the client and the endpoint to connect to. + * @param uri the WebSocket handshake URL + * @param client the WebSocket client + * @return the created instance + */ + public static WebSocketGraphQlTransport create(URI uri, WebSocketClient client) { + return new Builder(uri, client).build(); + } + + /** + * Return a builder with further options for the transport. + * @param uri the WebSocket handshake URL + * @param client the WebSocket client + * @return the builder instance + */ + public static Builder builder(URI uri, WebSocketClient client) { + return new Builder(uri, client); + } + + + /** + * Builder for {@link WebSocketGraphQlTransport} with an option to build a + * {@link GraphQlClient} instead, configured with the transport. + */ + public static final class Builder { + + private URI url; + + private final HttpHeaders headers = new HttpHeaders(); + + private WebSocketClient client; + + private CodecConfigurer codecsConfigurer = ClientCodecConfigurer.create(); + + @Nullable + private Object initPayload; + + private Consumer> connectionAckHandler = ackPayload -> {}; + + + Builder(URI url, WebSocketClient client) { + this.url = url; + this.client = client; + } + + + /** + * Set the URL for the WebSocket handshake request. + */ + public Builder url(URI uri) { + Assert.notNull(uri, "URI is required"); + this.url = uri; + return this; + } + + /** + * Add an HTTP header for the WebSocket handshake request. + */ + public Builder header(String name, String... values) { + Arrays.stream(values).forEach(value -> this.headers.add(name, value)); + return this; + } + + /** + * Provides access to every header declared so far with the possibility + * to add, replace, or remove. + */ + public Builder headers(Consumer headersConsumer) { + headersConsumer.accept(this.headers); + return this; + } + + /** + * Set the {@code WebSocketClient} to connect over. + */ + public Builder webSocketClient(WebSocketClient client) { + Assert.notNull(client, "WebSocketClient is required"); + this.client = client; + return this; + } + + /** + * Provide a {@code CodecConfigurer} that should contain encoders and + * decoders for JSON to be able to encode and decode GraphQL messages. + */ + public Builder codecConfigurer(CodecConfigurer codecConfigurer) { + this.codecsConfigurer = codecConfigurer; + return this; + } + + /** + * The payload to send with the "connection_init" message. + */ + public Builder connectionInitPayload(@Nullable Object connectionInitPayload) { + this.initPayload = connectionInitPayload; + return this; + } + + /** + * Handler for the payload received with the "connection_ack" message. + */ + public Builder connectionAckHandler(Consumer> ackHandler) { + this.connectionAckHandler = ackHandler; + return this; + } + + /** + * Build the transport instance. + */ + public WebSocketGraphQlTransport build() { + + Supplier mutateBuilder = () -> new Builder(this.url, this.client) + .headers(theHeaders -> theHeaders.putAll(this.headers)) + .codecConfigurer(codecsConfigurer.clone()) + .connectionInitPayload(this.initPayload) + .connectionAckHandler(this.connectionAckHandler); + + return new WebSocketGraphQlTransport(this.url, this.headers, this.client, + this.codecsConfigurer, this.initPayload, this.connectionAckHandler, mutateBuilder); + } + + /** + * Proceed to build a client that is configured with the transport from + * this builder. + */ + public GraphQlClient.Builder configureClient() { + return GraphQlClient.builder(build()); + } + + /** + * Build a client configured with the transport from this builder. + */ + public GraphQlClient buildClient() { + return GraphQlClient.builder(build()).build(); + } + + } + + + /** + * Client {@code WebSocketHandler} for GraphQL that deals with WebSocket + * concerns such as encoding and decoding of messages, {@link GraphQlSession} + * initialization and lifecycle, as well the lifecycle of the WebSocket. + * + *

This handler is for use as a singleton, but expects only one shared + * connection at a time. This is managed at a higher level by caching and + * re-using the {@link #getGraphQlSession() GraphQlSession} until closed at + * which point a new connection can be started. + */ + private static class GraphQlSessionHandler implements WebSocketHandler { + + private final WebSocketCodecDelegate codecDelegate; + + private final GraphQlWebSocketMessage connectionInitMessage; + + private final Consumer> connectionAckHandler; + + private Sinks.One graphQlSessionSink; + + private final AtomicBoolean stopped = new AtomicBoolean(); + + + GraphQlSessionHandler(CodecConfigurer codecConfigurer, + @Nullable Object connectionInitPayload, Consumer> connectionAckHandler) { + + this.codecDelegate = new WebSocketCodecDelegate(codecConfigurer); + this.connectionInitMessage = GraphQlWebSocketMessage.connectionInit(connectionInitPayload); + this.connectionAckHandler = connectionAckHandler; + this.graphQlSessionSink = Sinks.unsafe().one(); + } + + + @Override + public List getSubProtocols() { + return Collections.singletonList("graphql-transport-ws"); + } + + /** + * Return the {@link GraphQlSession} for sending requests. + * The {@code Mono} completes when the WebSocket session is connected and + * the "connection_init" and "connection_ack" messages are exchanged or + * returns an error if it fails for any reason. + */ + public Mono getGraphQlSession() { + return this.graphQlSessionSink.asMono(); + } + + /** + * When the handler is marked "stopped", i.e. set to {@code true}, new + * requests are rejected. When set to {@code true} they are allowed. + */ + public void setStopped(boolean stopped) { + this.stopped.set(stopped); + } + + /** + * Whether the handler is marked {@link #setStopped(boolean) "stopped"}. + */ + public boolean isStopped() { + return this.stopped.get(); + } + + + @Override + public Mono handle(WebSocketSession session) { + + Assert.state(sessionNotInitialized(), + "This handler supports only one session at a time, for shared use."); + + GraphQlSession graphQlSession = new GraphQlSession(session); + registerCloseStatusHandling(graphQlSession, session); + + Mono sendCompletion = + session.send(Flux.just(this.connectionInitMessage).concatWith(graphQlSession.getRequestFlux()) + .map(message -> this.codecDelegate.encode(session, message))); + + Mono receiveCompletion = session.receive() + .flatMap(webSocketMessage -> { + if (sessionNotInitialized()) { + try { + GraphQlWebSocketMessage message = this.codecDelegate.decode(webSocketMessage); + Assert.state(message.getType().equals("connection_ack"), + () -> "Unexpected message before connection_ack: " + message); + this.connectionAckHandler.accept(message.getPayload()); + if (logger.isDebugEnabled()) { + logger.debug(graphQlSession + " initialized"); + } + } + catch (Throwable ex) { + this.graphQlSessionSink.tryEmitError(ex); + return Mono.error(ex); + } + Sinks.EmitResult emitResult = this.graphQlSessionSink.tryEmitValue(graphQlSession); + if (emitResult.isFailure()) { + return Mono.error(new IllegalStateException( + "GraphQlSession initialized but could not be emitted: " + emitResult)); + } + } + else { + GraphQlWebSocketMessage message = this.codecDelegate.decode(webSocketMessage); + switch (message.getType()) { + case "next": + graphQlSession.handleNext(message); + break; + case "error": + graphQlSession.handleError(message); + break; + case "complete": + graphQlSession.handleComplete(message); + break; + default: + return Mono.error(new IllegalStateException("Unexpected message: " + message)); + } + } + return Mono.empty(); + }) + .then(); + + return Mono.zip(sendCompletion, receiveCompletion).then(); + } + + private boolean sessionNotInitialized() { + return !Boolean.TRUE.equals(this.graphQlSessionSink.scan(Scannable.Attr.TERMINATED)); + } + + private void registerCloseStatusHandling(GraphQlSession graphQlSession, WebSocketSession session) { + session.closeStatus() + .defaultIfEmpty(CloseStatus.NO_STATUS_CODE) + .doOnNext(closeStatus -> { + Exception ex = initDisconnectError(closeStatus, null, graphQlSession); + if (logger.isDebugEnabled()) { + logger.debug(ex.getMessage()); + } + graphQlSession.terminateRequests(ex); + }) + .doOnError(cause -> { + Exception ex = initDisconnectError(null, cause, graphQlSession); + if (logger.isErrorEnabled()) { + logger.error(ex.getMessage()); + } + graphQlSession.terminateRequests(ex); + }) + .doOnTerminate(() -> { + // Reset GraphQlSession sink to be ready to connect again + this.graphQlSessionSink = Sinks.unsafe().one(); + }) + .subscribe(); + } + + private Exception initDisconnectError( + @Nullable CloseStatus status, @Nullable Throwable ex, GraphQlSession graphQlSession) { + + String reason = graphQlSession + " disconnected"; + if (isStopped()) { + reason = graphQlSession + " was stopped"; + } + else if (ex != null) { + reason += ", closeStatus() completed with error " + ex; + } + else if (status != null && !status.equals(CloseStatus.NO_STATUS_CODE)) { + reason += " with " + status; + } + else { + reason += " without a status"; + } + return new IllegalStateException(reason); + } + + /** + * This must be called from code that calls the {@code WebSocketClient} + * when execution completes with an error, which includes connection and + * session handling issues, and handler is unaware of connection issues + * otherwise. + * + *

The exception is logged which may provide further information, + * beyond the CloseStatus, when closed locally due to handling completing + * with an error. The error is routed to subscribers of + * {@link #getGraphQlSession()} which is necessary for connection issues. + */ + public void handleWebSocketSessionError(Throwable ex) { + + if (logger.isDebugEnabled()) { + logger.debug("Session handling error: " + ex.getMessage(), ex); + } + else if (logger.isErrorEnabled()) { + logger.error("Session handling error: " + ex.getMessage()); + } + + this.graphQlSessionSink.tryEmitError(ex); + } + + } + + + /** + * Session that deals with GraphQL level concerns such as sending requests, + * handling and routing responses, managing the lifecycle of streams, and + * to allow higher level code to be notified of closing or to close the + * underlying WebSocket. + */ + private static class GraphQlSession { + + private final DisposableConnection connection; + + private final AtomicLong requestIndex = new AtomicLong(); + + private final Sinks.Many requestSink = Sinks.many().unicast().onBackpressureBuffer(); + + private final Map> resultSinks = new ConcurrentHashMap<>(); + + private final Map> streamingSinks = new ConcurrentHashMap<>(); + + + GraphQlSession(WebSocketSession webSocketSession) { + this.connection = DisposableConnection.from(webSocketSession); + } + + + /** + * Return the {@code Flux} of GraphQL requests to send as WebSocket messages. + */ + public Flux getRequestFlux() { + return this.requestSink.asFlux(); + } + + public Mono execute(RequestInput requestInput) { + String id = String.valueOf(this.requestIndex.incrementAndGet()); + try { + GraphQlWebSocketMessage message = GraphQlWebSocketMessage.subscribe(id, requestInput); + Sinks.One sink = Sinks.one(); + this.resultSinks.put(id, sink); + trySend(message); + return sink.asMono().doOnCancel(() -> this.resultSinks.remove(id)); + } + catch (Exception ex) { + this.resultSinks.remove(id); + return Mono.error(ex); + } + } + + public Flux executeSubscription(RequestInput requestInput) { + String id = String.valueOf(this.requestIndex.incrementAndGet()); + try { + GraphQlWebSocketMessage message = GraphQlWebSocketMessage.subscribe(id, requestInput); + Sinks.Many sink = Sinks.many().unicast().onBackpressureBuffer(); + this.streamingSinks.put(id, sink); + trySend(message); + return sink.asFlux().doOnCancel(() -> cancelStream(id)); + } + catch (Exception ex) { + this.streamingSinks.remove(id); + return Flux.error(ex); + } + } + + // TODO: queue to serialize sending? + + private void trySend(GraphQlWebSocketMessage message) { + Sinks.EmitResult emitResult = null; + for (int i = 0; i < 100; i++) { + emitResult = this.requestSink.tryEmitNext(message); + if (emitResult != Sinks.EmitResult.FAIL_NON_SERIALIZED) { + break; + } + } + Assert.state(emitResult.isSuccess(), "Failed to send request: " + emitResult); + } + + private void cancelStream(String id) { + Sinks.Many streamSink = this.streamingSinks.remove(id); + if (streamSink != null) { + try { + trySend(GraphQlWebSocketMessage.complete(id)); + } + catch (Exception ex) { + if (logger.isErrorEnabled()) { + logger.error("Closing " + this.connection.getDescription() + + " after failure to send 'complete' for subscription id='" + id + "'."); + } + this.connection.close().subscribe(); + } + } + } + + /** + * Handle a "next" message and route to its recipient. + */ + public void handleNext(GraphQlWebSocketMessage message) { + String id = message.getId(); + Sinks.One sink = this.resultSinks.remove(id); + Sinks.Many streamingSink = this.streamingSinks.get(id); + + if (sink == null && streamingSink == null) { + if (logger.isDebugEnabled()) { + logger.debug("No receiver for message: " + message); + } + return; + } + + ExecutionResult result = new MapExecutionResult(message.getPayload()); + Sinks.EmitResult emitResult = (sink != null ? sink.tryEmitValue(result) : streamingSink.tryEmitNext(result)); + if (emitResult.isFailure()) { + // Just log: cannot overflow, is serialized, and cancel is handled in doOnCancel + if (logger.isDebugEnabled()) { + logger.debug("Message: " + message + " could not be emitted: " + emitResult); + } + } + } + + /** + * Handle an "error" message, turning it into an {@link ExecutionResult} + * for a single result response, or signaling an error to streams. + */ + public void handleError(GraphQlWebSocketMessage message) { + String id = message.getId(); + Sinks.One sink = this.resultSinks.remove(id); + Sinks.Many streamingSink = this.streamingSinks.remove(id); + + if (sink == null && streamingSink == null ) { + if (logger.isDebugEnabled()) { + logger.debug("No receiver for message: " + message); + } + return; + } + + List> payload = message.getPayload(); + + Sinks.EmitResult emitResult; + if (sink != null) { + ExecutionResult result = new MapExecutionResult(Collections.singletonMap("errors", payload)); + emitResult = sink.tryEmitValue(result); + } + else { + List graphQLErrors = MapGraphQlError.fromMapList(payload); + Exception ex = new SubscriptionErrorException(graphQLErrors); + emitResult = streamingSink.tryEmitError(ex); + } + + if (emitResult.isFailure() && logger.isDebugEnabled()) { + logger.debug("Error: " + message + " could not be emitted: " + emitResult); + } + } + + /** + * Handle a "complete" message. + */ + public void handleComplete(GraphQlWebSocketMessage message) { + Sinks.One resultSink = this.resultSinks.remove(message.getId()); + Sinks.Many streamingResultSink = this.streamingSinks.remove(message.getId()); + + if (resultSink != null) { + resultSink.tryEmitEmpty(); + } + else if (streamingResultSink != null) { + streamingResultSink.tryEmitComplete(); + } + } + + /** + * Return a {@code Mono} that completes when the connection is closed + * for any reason. + */ + public Mono notifyWhenClosed() { + return this.connection.notifyWhenClosed(); + } + + /** + * Close the underlying connection. + */ + public Mono close() { + return this.connection.close(); + } + + /** + * Terminate and clean all in-progress requests with the given error. + */ + public void terminateRequests(Exception ex) { + this.resultSinks.values().forEach(sink -> sink.tryEmitError(ex)); + this.streamingSinks.values().forEach(sink -> sink.tryEmitError(ex)); + this.resultSinks.clear(); + this.streamingSinks.clear(); + } + + @Override + public String toString() { + return "GraphQlSession over " + this.connection.getDescription(); + } + + } + + + /** + * Minimal abstraction to decouple the {@link GraphQlSession} from the + * underlying {@code WebSocketSession}. + */ + private interface DisposableConnection { + + Mono close(); + + Mono notifyWhenClosed(); + + String getDescription(); + + + static DisposableConnection from(WebSocketSession session) { + + return new DisposableConnection() { + + @Override + public Mono close() { + return session.close(); + } + + @Override + public Mono notifyWhenClosed() { + return session.closeStatus().then(); + } + + @Override + public String getDescription() { + return session.toString(); + } + + }; + } + } + + +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java b/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java new file mode 100644 index 00000000..a293d513 --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java @@ -0,0 +1,332 @@ +/* + * Copyright 2002-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.graphql.client; + +import java.io.IOException; +import java.net.URI; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; + +import graphql.ExecutionResult; +import graphql.GraphQLError; +import graphql.GraphqlErrorBuilder; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import org.springframework.graphql.RequestInput; +import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage; +import org.springframework.http.HttpHeaders; +import org.springframework.web.reactive.socket.CloseStatus; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.WebSocketSession; +import org.springframework.web.reactive.socket.client.WebSocketClient; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link WebSocketGraphQlTransport}. + * @author Rossen Stoyanchev + */ +public class MockWebSocketGraphQlTransportTests { + + private final static Duration TIMEOUT = Duration.ofSeconds(5); + + private static final WebSocketCodecDelegate CODEC_DELEGATE = new WebSocketCodecDelegate(); + + + private final MockWebSocketServer mockServer = new MockWebSocketServer(); + + private final TestWebSocketClient testClient = new TestWebSocketClient(this.mockServer); + + private final WebSocketGraphQlTransport transport = + WebSocketGraphQlTransport.builder(URI.create("/"), this.testClient).build(); + + private final ExecutionResult result1 = MapExecutionResult.forData(Collections.singletonMap("key1", "value1")); + + private final ExecutionResult result2 = MapExecutionResult.forData(Collections.singletonMap("key2", "value2")); + + + @Test + void request() { + RequestInput input = this.mockServer.expectOperation("Query1").andRespond(this.result1); + + StepVerifier.create(this.transport.execute(input)) + .expectNext(this.result1).expectComplete() + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input)); + } + + @Test + void requestStream() { + RequestInput input = this.mockServer.expectOperation("Sub1").andStream(Flux.just(this.result1, result2)); + + StepVerifier.create(this.transport.executeSubscription(input)) + .expectNext(this.result1, result2).expectComplete() + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input)); + } + + @Test + void requestError() { + RequestInput input = this.mockServer.expectOperation("Query1") + .andRespondWithError(GraphqlErrorBuilder.newError().message("boo").build()); + + StepVerifier.create(this.transport.execute(input)) + .consumeNextWith(result -> { + assertThat(result.isDataPresent()).isFalse(); + assertThat(result.getErrors()).extracting(GraphQLError::getMessage).containsExactly("boo"); + }) + .expectComplete() + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input)); + } + + @Test + void requestStreamError() { + RequestInput input = this.mockServer.expectOperation("Sub1") + .andStreamWithError(Flux.just(this.result1), GraphqlErrorBuilder.newError().message("boo").build()); + + StepVerifier.create(this.transport.executeSubscription(input)) + .expectNext(this.result1) + .expectErrorSatisfies(actualEx -> { + List errorList = ((SubscriptionErrorException) actualEx).getErrors(); + assertThat(errorList).extracting(GraphQLError::getMessage).containsExactly("boo"); + }) + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input)); + } + + @Test + void requestCancelled() { + RequestInput input = this.mockServer.expectOperation("Query1").andRespond(Mono.never()); + + StepVerifier.create(this.transport.execute(input)) + .thenAwait(Duration.ofMillis(200)) + .thenCancel() + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input)); + } + + @Test + void requestStreamCancelled() { + RequestInput input = this.mockServer.expectOperation("s1") + .andStream(Flux.just(this.result1).concatWith(Flux.never())); + + StepVerifier.create(this.transport.executeSubscription(input)) + .expectNext(this.result1) + .thenAwait(Duration.ofMillis(200)) + .thenCancel() + .verify(TIMEOUT); + + assertActualClientMessages( + GraphQlWebSocketMessage.connectionInit(null), + GraphQlWebSocketMessage.subscribe("1", input), + GraphQlWebSocketMessage.complete("1")); + } + + @Test + void start() { + MockWebSocketServer handler = new MockWebSocketServer(); + handler.connectionInitHandler(payload -> Mono.just(Collections.singletonMap("key", payload.get("key") + "Ack"))); + + TestWebSocketClient client = new TestWebSocketClient(handler); + Map initPayload = Collections.singletonMap("key", "valueInit"); + AtomicReference> connectionAckRef = new AtomicReference<>(); + + WebSocketGraphQlTransport transport = WebSocketGraphQlTransport.builder(URI.create("/"), client) + .connectionInitPayload(initPayload) + .connectionAckHandler(connectionAckRef::set) + .build(); + + transport.start().block(TIMEOUT); + + assertThat(client.getConnection(0).isOpen()).isTrue(); + assertThat(connectionAckRef.get()).isEqualTo(Collections.singletonMap("key", "valueInitAck")); + assertActualClientMessages(client.getConnection(0), GraphQlWebSocketMessage.connectionInit(initPayload)); + } + + @Test + void stop() { + + // Start + this.transport.start().block(TIMEOUT); + assertThat(this.testClient.getConnectionCount()).isEqualTo(1); + assertThat(this.testClient.getConnection(0).isOpen()).isTrue(); + + // Stop + this.transport.stop().block(TIMEOUT); + assertThat(this.testClient.getConnection(0).isOpen()).isFalse(); + assertThat(this.testClient.getConnection(0).closeStatus().block(TIMEOUT)).isEqualTo(CloseStatus.NORMAL); + + // New requests are rejected + RequestInput input = this.mockServer.expectOperation("Query1").andRespond(this.result1); + StepVerifier.create(this.transport.execute(input)) + .expectErrorMessage("WebSocketGraphQlTransport has been stopped") + .verify(TIMEOUT); + + // Start + this.transport.start().block(TIMEOUT); + assertThat(this.testClient.getConnectionCount()).isEqualTo(2); + assertThat(this.testClient.getConnection(1).isOpen()).isTrue(); + + // Requests allowed again + input = this.mockServer.expectOperation("Query1").andRespond(this.result1); + StepVerifier.create(this.transport.execute(input)) + .expectNext(this.result1).expectComplete() + .verify(TIMEOUT); + } + + @Test + void sessionIsCachedUntilClosed() { + + RequestInput input1 = this.mockServer.expectOperation("Query1").andRespond(this.result1); + StepVerifier.create(this.transport.execute(input1)).expectNext(this.result1).expectComplete().verify(TIMEOUT); + + assertThat(this.testClient.getConnectionCount()).isEqualTo(1); + TestWebSocketConnection originalConnection = this.testClient.getConnection(0); + + RequestInput input2 = this.mockServer.expectOperation("Query2").andRespond(this.result2); + StepVerifier.create(this.transport.execute(input2)).expectNext(this.result2).expectComplete().verify(TIMEOUT); + + assertThat(this.testClient.getConnectionCount()).isEqualTo(1); + assertThat(this.testClient.getConnection(0)).isSameAs(originalConnection); + + // Close the connection + originalConnection.closeServerSession(CloseStatus.NORMAL).block(TIMEOUT); + + input1 = this.mockServer.expectOperation("Query1").andRespond(this.result1); + StepVerifier.create(this.transport.execute(input1)).expectNext(this.result1).expectComplete().verify(TIMEOUT); + + assertThat(this.testClient.getConnectionCount()).isEqualTo(2); + assertThat(this.testClient.getConnection(1)).isNotSameAs(originalConnection); + } + + @Test + void errorOnConnect() { + + // Connection errors should be routed, no hanging on start + + IOException ex = new IOException("Connect failure"); + + WebSocketClient client = mock(WebSocketClient.class); + when(client.execute(any(URI.class), any(HttpHeaders.class), any(WebSocketHandler.class))).thenReturn(Mono.error(ex)); + + StepVerifier.create(createTransport(client).start()) + .expectErrorMessage(ex.getMessage()) + .verify(TIMEOUT); + } + + @Test + void errorBeforeConnectionAck() { + + // Errors before GraphQL session initialized should be routed, no hanging on start + + MockWebSocketServer handler = new MockWebSocketServer(); + handler.connectionInitHandler(initPayload -> Mono.error(new IllegalStateException("boo"))); + + TestWebSocketClient client = new TestWebSocketClient(handler); + + StepVerifier.create(createTransport(client).start()) + .expectErrorMessage("boo") + .verify(TIMEOUT); + } + + @Test + void errorDuringResponseHandling() { + + // Response handling errors that close the connection should terminate outstanding requests + + TestWebSocketClient client = new TestWebSocketClient(new UnexpectedResponseHandler()); + WebSocketGraphQlTransport transport = createTransport(client); + + String expectedMessage = "GraphQlSession over client-session-1 disconnected " + + "with CloseStatus[code=1002, reason=null]"; + + StepVerifier.create(transport.execute(new RequestInput("Query1", null, null, null, ""))) + .expectErrorMessage(expectedMessage) + .verify(TIMEOUT); + } + + private WebSocketGraphQlTransport createTransport(WebSocketClient client) { + return WebSocketGraphQlTransport.builder(URI.create("/"), client).build(); + } + + private void assertActualClientMessages(GraphQlWebSocketMessage... expectedMessages) { + assertActualClientMessages(this.testClient.getConnection(0), expectedMessages); + } + + private void assertActualClientMessages( + TestWebSocketConnection connection, GraphQlWebSocketMessage... expectedMessages) { + + List actualMessages = connection.getClientMessages().stream() + .map(CODEC_DELEGATE::decode) + .collect(Collectors.toList()); + + assertThat(actualMessages).containsExactly(expectedMessages); + } + + + /** + * Server handler that returns an unexpected (client) message. + */ + private static class UnexpectedResponseHandler implements WebSocketHandler { + + private final WebSocketCodecDelegate codecDelegate = new WebSocketCodecDelegate(); + + + @Override + public Mono handle(WebSocketSession session) { + return session.send(session.receive().flatMap(webSocketMessage -> { + + GraphQlWebSocketMessage requestMessage = this.codecDelegate.decode(webSocketMessage); + String id = requestMessage.getId(); + + GraphQlWebSocketMessage responseMessage = (requestMessage.getType().equals("connection_init") ? + GraphQlWebSocketMessage.connectionAck(null) : + GraphQlWebSocketMessage.subscribe(id, new RequestInput("..", null, null, null, ""))); + + return Flux.just(this.codecDelegate.encode(session, responseMessage)); + })); + } + + } + +} diff --git a/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketServer.java b/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketServer.java new file mode 100644 index 00000000..276428b1 --- /dev/null +++ b/spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketServer.java @@ -0,0 +1,200 @@ +/* + * Copyright 2002-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.graphql.client; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.function.Function; + +import graphql.ExecutionResult; +import graphql.GraphQLError; +import org.reactivestreams.Publisher; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import org.springframework.graphql.RequestInput; +import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage; +import org.springframework.lang.Nullable; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.WebSocketSession; + +/** + * GraphQL over WebSocket handler to use as a server-side + * {@link WebSocketHandler} that is configured with expected requests and + * the responses to send. + * + * @author Rossen Stoyanchev + */ +public class MockWebSocketServer implements WebSocketHandler { + + @Nullable + private Function, Mono> connectionInitHandler; + + private final Map, Exchange> expectedExchanges = new LinkedHashMap<>(); + + private final WebSocketCodecDelegate codecDelegate = new WebSocketCodecDelegate(); + + + /** + * Configure a handler for the "connection_init" message. + * @param handler accepts "connection_init" and returns the "connection_ack" payload + */ + public void connectionInitHandler(Function, Mono> handler) { + this.connectionInitHandler = handler; + } + + /** + * Add the GraphQL operation for an expected request and then specify the + * response to send back. + */ + public ResponseSpec expectOperation(String operation) { + Exchange exchange = new Exchange(operation); + this.expectedExchanges.put(exchange.getInput().toMap(), exchange); + return exchange; + } + + + @Override + public Mono handle(WebSocketSession session) { + return session.send(session.receive() + .map(codecDelegate::decode) + .flatMap(this::handleMessage) + .map(message -> codecDelegate.encode(session, message))); + } + + @SuppressWarnings("SuspiciousMethodCalls") + private Publisher handleMessage(GraphQlWebSocketMessage message) { + if ("connection_init".equals(message.getType())) { + if (this.connectionInitHandler == null) { + return Flux.just(GraphQlWebSocketMessage.connectionAck(null)); + } + else { + Map payload = message.getPayload(); + return this.connectionInitHandler.apply(payload).map(GraphQlWebSocketMessage::connectionAck); + } + } + if ("subscribe".equals(message.getType())) { + String id = message.getId(); + Exchange request = expectedExchanges.get(message.getPayload()); + if (id == null || request == null) { + return Flux.error(new IllegalStateException("Unexpected request: " + message)); + } + return request.getResponseFlux() + .map(result -> GraphQlWebSocketMessage.next(id, result)) + .concatWithValues( + request.getError() != null ? + GraphQlWebSocketMessage.error(id, request.getError()) : + GraphQlWebSocketMessage.complete(id)); + } + if ("complete".equals(message.getType())) { + return Flux.empty(); + } + return Flux.error(new IllegalStateException("Unexpected message: " + message)); + } + + + public interface ResponseSpec { + + /** + * Respond with the given a single result. + */ + RequestInput andRespond(ExecutionResult result); + + /** + * Respond with the given a single result {@code Mono}. + */ + RequestInput andRespond(Mono resultMono); + + /** + * Respond with a GraphQL over WebSocket "error" message. + */ + RequestInput andRespondWithError(GraphQLError error); + + /** + * Respond with the given stream of responses. + */ + RequestInput andStream(Flux resultFlux); + + /** + * Respond with the given stream of responses and terminate with an error. + */ + RequestInput andStreamWithError(Flux resultFlux, GraphQLError error); + + } + + + private static class Exchange implements ResponseSpec { + + private final RequestInput requestInput; + + private Flux responseFlux = Flux.empty(); + + @Nullable + private GraphQLError error; + + + private Exchange(String operation) { + this.requestInput = new RequestInput(operation, null, null, null, ""); + } + + @Override + public RequestInput andRespond(ExecutionResult result) { + return addResponse(Flux.just(result), null); + } + + @Override + public RequestInput andRespond(Mono resultMono) { + return addResponse(Flux.from(resultMono), null); + } + + @Override + public RequestInput andRespondWithError(GraphQLError error) { + return addResponse(Flux.empty(), error); + } + + @Override + public RequestInput andStream(Flux resultFlux) { + return addResponse(resultFlux, null); + } + + @Override + public RequestInput andStreamWithError(Flux resultFlux, GraphQLError error) { + return addResponse(resultFlux, error); + } + + private RequestInput addResponse(Flux resultFlux, @Nullable GraphQLError error) { + this.responseFlux = resultFlux; + this.error = error; + return this.requestInput; + } + + public RequestInput getInput() { + return this.requestInput; + } + + public Flux getResponseFlux() { + return this.responseFlux; + } + + @Nullable + public GraphQLError getError() { + return this.error; + } + + } + +}