From 5a77b9367170c576942d4dc5b5f0c1faa87ba857 Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Mon, 18 Jan 2021 21:18:18 +0000 Subject: [PATCH] Minor refactoring in GraphQLWebSocketHandler The subscriptions map is now per session and not global. The rest is polishing. Minor refactoring --- .../webflux/GraphQLWebSocketHandler.java | 92 ++++++++++--------- 1 file changed, 47 insertions(+), 45 deletions(-) diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java index 6abbb308..2e311475 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java @@ -61,6 +61,7 @@ import org.springframework.web.reactive.socket.WebSocketSession; /** * WebSocketHandler for GraphQL based on * GraphQL Over WebSocket Protocol + * and for use in a Spring WebFlux application. */ public class GraphQLWebSocketHandler implements WebSocketHandler { @@ -74,35 +75,28 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { private final WebInterceptorExecutionChain executionChain; - private final Duration initTimeoutDuration; - private final Decoder decoder; private final Encoder encoder; - private final Map subscriptions = new ConcurrentHashMap<>(); + private final Duration initTimeoutDuration; - @Override - public List getSubProtocols() { - return SUB_PROTOCOL_LIST; - } - /** * Create a new instance. * @param graphQL the GraphQL instance to use for query execution * @param interceptors 0 or more interceptors to customize input and output * @param configurer codec configurer for JSON encoding and decoding - * @param initTimeoutDuration the time within which the -* {@code CONNECTION_INIT} type message must be received. + * @param initTimeoutDuration the time within which the {@code CONNECTION_INIT} + * type message must be received. */ public GraphQLWebSocketHandler(GraphQL graphQL, List interceptors, ServerCodecConfigurer configurer, Duration initTimeoutDuration) { this.executionChain = new WebInterceptorExecutionChain(graphQL, interceptors); - this.initTimeoutDuration = initTimeoutDuration; this.decoder = initDecoder(configurer); this.encoder = initEncoder(configurer); + this.initTimeoutDuration = initTimeoutDuration; } private static Decoder initDecoder(ServerCodecConfigurer configurer) { @@ -122,13 +116,23 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { } + @Override + public List getSubProtocols() { + return SUB_PROTOCOL_LIST; + } + + @Override public Mono handle(WebSocketSession session) { - AtomicBoolean initialized = new AtomicBoolean(); + + // Session state + AtomicBoolean connectionInitProcessed = new AtomicBoolean(); + Map subscriptions = new ConcurrentHashMap<>(); + Mono.delay(this.initTimeoutDuration) .then(Mono.defer(() -> - initialized.compareAndSet(false, true) ? - GraphQLStatus.initTimeout(session) : + connectionInitProcessed.compareAndSet(false, true) ? + GraphQLStatus.closeWithInitTimeout(session) : Mono.empty())) .subscribe(); @@ -138,38 +142,38 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { String id = (String) map.get("id"); MessageType messageType = MessageType.resolve((String) map.get("type")); if (messageType == null) { - return GraphQLStatus.invalidMessage(session); + return GraphQLStatus.closeWithInvalidMessage(session); } switch (messageType) { case SUBSCRIBE: - if (!initialized.get()) { - return GraphQLStatus.unauthorized(session); + if (!connectionInitProcessed.get()) { + return GraphQLStatus.closeWithUnauthorized(session); } if (id == null) { - return GraphQLStatus.invalidMessage(session); + return GraphQLStatus.closeWithInvalidMessage(session); } HandshakeInfo handshakeInfo = session.getHandshakeInfo(); WebSocketInput input = new WebSocketInput(handshakeInfo, id, getPayload(map)); if (logger.isDebugEnabled()) { logger.debug("Executing: " + input); } - return executionChain.execute(input) - .flatMapMany(output -> handleWebOutput(session, input.id(), output)); + return executionChain.execute(input).flatMapMany(output -> + handleWebOutput(session, input.id(), subscriptions, output)); case COMPLETE: if (id != null) { - Subscription subscription = this.subscriptions.remove(id); + Subscription subscription = subscriptions.remove(id); if (subscription != null) { subscription.cancel(); } } return Flux.empty(); case CONNECTION_INIT: - if (!initialized.compareAndSet(false, true)) { - return GraphQLStatus.tooManyInitRequests(session); + if (!connectionInitProcessed.compareAndSet(false, true)) { + return GraphQLStatus.closeWithTooManyInitRequests(session); } return Flux.just(encode(session, null, MessageType.CONNECTION_ACK, null)); default: - return GraphQLStatus.invalidMessage(session); + return GraphQLStatus.closeWithInvalidMessage(session); } }); @@ -190,7 +194,9 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { } @SuppressWarnings("unchecked") - private Flux handleWebOutput(WebSocketSession session, String id, WebOutput output) { + private Flux handleWebOutput( + WebSocketSession session, String id, Map subscriptions, WebOutput output) { + if (logger.isDebugEnabled()) { logger.debug("Execution result ready" + (!CollectionUtils.isEmpty(output.getErrors()) ? @@ -202,7 +208,7 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { // Subscription outputFlux = Flux.from((Publisher) output.getData()) .doOnSubscribe(subscription -> { - Subscription previous = this.subscriptions.putIfAbsent(id, subscription); + Subscription previous = subscriptions.putIfAbsent(id, subscription); if (previous != null) { throw new SubscriptionExistsException(); } @@ -223,7 +229,7 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { .concatWith(Mono.defer(() -> Mono.just(encode(session, id, MessageType.COMPLETE, null)))) .onErrorResume(ex -> { if (ex instanceof SubscriptionExistsException) { - return GraphQLStatus.subscriptionExists(session, id); + return GraphQLStatus.closeWithSubscriptionExists(session, id); } ErrorType errorType = ErrorType.DataFetchingException; String message = ex.getMessage(); @@ -295,40 +301,36 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { private static class GraphQLStatus { - private static final CloseStatus INVALID_MESSAGE_STATUS = - new CloseStatus(4400, "Invalid message"); + private static final CloseStatus INVALID_MESSAGE_STATUS = new CloseStatus(4400, "Invalid message"); - private static final CloseStatus UNAUTHORIZED_STATUS = - new CloseStatus(4401, "Unauthorized"); + private static final CloseStatus UNAUTHORIZED_STATUS = new CloseStatus(4401, "Unauthorized"); - private static final CloseStatus INIT_TIMEOUT_STATUS = - new CloseStatus(4408, "Connection initialisation timeout"); + private static final CloseStatus INIT_TIMEOUT_STATUS = new CloseStatus(4408, "Connection initialisation timeout"); - private static final CloseStatus TOO_MANY_INIT_REQUESTS_STATUS = - new CloseStatus(4429, "Too many initialisation requests"); + private static final CloseStatus TOO_MANY_INIT_REQUESTS_STATUS = new CloseStatus(4429, "Too many initialisation requests"); - public static Flux invalidMessage(WebSocketSession session) { - return closeSession(session, INVALID_MESSAGE_STATUS); + public static Flux closeWithInvalidMessage(WebSocketSession session) { + return closeInternal(session, INVALID_MESSAGE_STATUS); } - public static Flux unauthorized(WebSocketSession session) { - return closeSession(session, UNAUTHORIZED_STATUS); + public static Flux closeWithUnauthorized(WebSocketSession session) { + return closeInternal(session, UNAUTHORIZED_STATUS); } - public static Mono initTimeout(WebSocketSession session) { + public static Mono closeWithInitTimeout(WebSocketSession session) { return session.close(INIT_TIMEOUT_STATUS); } - public static Flux subscriptionExists(WebSocketSession session, String id) { - return closeSession(session, new CloseStatus(4409, "Subscriber for " + id + " already exists")); + public static Flux closeWithSubscriptionExists(WebSocketSession session, String id) { + return closeInternal(session, new CloseStatus(4409, "Subscriber for " + id + " already exists")); } - public static Flux tooManyInitRequests(WebSocketSession session) { - return closeSession(session, TOO_MANY_INIT_REQUESTS_STATUS); + public static Flux closeWithTooManyInitRequests(WebSocketSession session) { + return closeInternal(session, TOO_MANY_INIT_REQUESTS_STATUS); } - private static Flux closeSession(WebSocketSession session, CloseStatus status) { + private static Flux closeInternal(WebSocketSession session, CloseStatus status) { return session.close(status).thenMany(Mono.empty()); } }