WebSocket handlers support keepalive PING messages

Closes gh-534

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/server/webflux/GraphQlWebSocketHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -44,6 +44,7 @@
 import org.springframework.graphql.server.support.GraphQlWebSocketMessage;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.codec.CodecConfigurer;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.web.reactive.socket.CloseStatus;
@@ -72,10 +73,13 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 
 	private final WebSocketGraphQlInterceptor webSocketInterceptor;
 
-	private final WebSocketCodecDelegate webSocketCodecDelegate;
+	private final WebSocketCodecDelegate codecDelegate;
 
 	private final Duration initTimeoutDuration;
 
+	@Nullable
+	private final Duration keepAliveDuration;
+
 
 	/**
 	 * Create a new instance.
@@ -87,12 +91,30 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 	public GraphQlWebSocketHandler(
 			WebGraphQlHandler graphQlHandler, CodecConfigurer codecConfigurer, Duration connectionInitTimeout) {
 
+		this(graphQlHandler, codecConfigurer, connectionInitTimeout, null);
+	}
+
+	/**
+	 * Create a new instance.
+	 * @param graphQlHandler common handler for GraphQL over WebSocket requests
+	 * @param codecConfigurer codec configurer for JSON encoding and decoding
+	 * @param connectionInitTimeout how long to wait after the establishment of
+	 * the WebSocket for the {@code "connection_ini"} message from the client.
+	 * @param keepAliveDuration how frequently to send ping messages; if not
+	 * set then ping messages are not sent.
+	 * @since 1.3
+	 */
+	public GraphQlWebSocketHandler(
+			WebGraphQlHandler graphQlHandler, CodecConfigurer codecConfigurer,
+			Duration connectionInitTimeout, @Nullable Duration keepAliveDuration) {
+
 		Assert.notNull(graphQlHandler, "WebGraphQlHandler is required");
 
 		this.graphQlHandler = graphQlHandler;
 		this.webSocketInterceptor = this.graphQlHandler.getWebSocketInterceptor();
-		this.webSocketCodecDelegate = new WebSocketCodecDelegate(codecConfigurer);
+		this.codecDelegate = new WebSocketCodecDelegate(codecConfigurer);
 		this.initTimeoutDuration = connectionInitTimeout;
+		this.keepAliveDuration = keepAliveDuration;
 	}
 
 
@@ -137,7 +159,7 @@ public Mono<Void> handle(WebSocketSession session) {
 				.subscribe();
 
 		return session.send(session.receive().flatMap((webSocketMessage) -> {
-			GraphQlWebSocketMessage message = this.webSocketCodecDelegate.decode(webSocketMessage);
+			GraphQlWebSocketMessage message = this.codecDelegate.decode(webSocketMessage);
 			String id = message.getId();
 			Map<String, Object> payload = message.getPayload();
 			switch (message.resolvedType()) {
@@ -159,7 +181,7 @@ public Mono<Void> handle(WebSocketSession session) {
 							.doOnTerminate(() -> subscriptions.remove(id));
 				}
 				case PING -> {
-					return Flux.just(this.webSocketCodecDelegate.encode(session, GraphQlWebSocketMessage.pong(null)));
+					return Flux.just(this.codecDelegate.encode(session, GraphQlWebSocketMessage.pong(null)));
 				}
 				case COMPLETE -> {
 					if (id != null) {
@@ -176,11 +198,16 @@ public Mono<Void> handle(WebSocketSession session) {
 					if (!connectionInitPayloadRef.compareAndSet(null, payload)) {
 						return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
 					}
-					return this.webSocketInterceptor.handleConnectionInitialization(sessionInfo, payload)
+					Flux<WebSocketMessage> flux = this.webSocketInterceptor.handleConnectionInitialization(sessionInfo, payload)
 							.defaultIfEmpty(Collections.emptyMap())
-							.map((ackPayload) -> this.webSocketCodecDelegate.encodeConnectionAck(session, ackPayload))
-							.flux()
-							.onErrorResume((ex) -> GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS));
+							.map((ackPayload) -> this.codecDelegate.encodeConnectionAck(session, ackPayload))
+							.flux();
+					if (this.keepAliveDuration != null) {
+						flux = flux.mergeWith(Flux.interval(this.keepAliveDuration, this.keepAliveDuration)
+								.filter((aLong) -> !this.codecDelegate.checkMessagesEncodedAndClear())
+								.map((aLong) -> this.codecDelegate.encode(session, GraphQlWebSocketMessage.ping(null))));
+					}
+					return flux.onErrorResume((ex) -> GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS));
 				}
 				default -> {
 					return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
@@ -218,14 +245,14 @@ private Flux<WebSocketMessage> handleResponse(WebSocketSession session, String i
 		}
 
 		return responseFlux
-				.map((responseMap) -> this.webSocketCodecDelegate.encodeNext(session, id, responseMap))
-				.concatWith(Mono.fromCallable(() -> this.webSocketCodecDelegate.encodeComplete(session, id)))
+				.map((responseMap) -> this.codecDelegate.encodeNext(session, id, responseMap))
+				.concatWith(Mono.fromCallable(() -> this.codecDelegate.encodeComplete(session, id)))
 				.onErrorResume((ex) -> {
 					if (ex instanceof SubscriptionExistsException) {
 						CloseStatus status = new CloseStatus(4409, "Subscriber for " + id + " already exists");
 						return GraphQlStatus.close(session, status);
 					}
-					return Mono.fromCallable(() -> this.webSocketCodecDelegate.encodeError(session, id, ex));
+					return Mono.fromCallable(() -> this.codecDelegate.encodeError(session, id, ex));
 				});
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/server/webflux/WebSocketCodecDelegate.java

@@ -54,6 +54,8 @@ final class WebSocketCodecDelegate {
 
 	private final Encoder<?> encoder;
 
+	private boolean messagesEncoded;
+
 
 	WebSocketCodecDelegate(CodecConfigurer codecConfigurer) {
 		Assert.notNull(codecConfigurer, "CodecConfigurer is required");
@@ -84,6 +86,8 @@ <T> WebSocketMessage encode(WebSocketSession session, GraphQlWebSocketMessage me
 		DataBuffer buffer = ((Encoder<T>) this.encoder).encodeValue(
 				(T) message, session.bufferFactory(), MESSAGE_TYPE, MimeTypeUtils.APPLICATION_JSON, null);
 
+		this.messagesEncoded = true;
+
 		return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer);
 	}
 
@@ -115,4 +119,10 @@ WebSocketMessage encodeComplete(WebSocketSession session, String id) {
 		return encode(session, GraphQlWebSocketMessage.complete(id));
 	}
 
+	boolean checkMessagesEncodedAndClear() {
+		boolean result = this.messagesEncoded;
+		this.messagesEncoded = false;
+		return result;
+	}
+
 }

spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java

@@ -104,8 +104,12 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 
 	private final HttpMessageConverter<?> converter;
 
+	@Nullable
+	private final Duration keepAliveDuration;
+
 	private final Map<String, SessionState> sessionInfoMap = new ConcurrentHashMap<>();
 
+
 	/**
 	 * Create a new instance.
 	 * @param graphQlHandler common handler for GraphQL over WebSocket requests
@@ -116,6 +120,23 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 	public GraphQlWebSocketHandler(
 			WebGraphQlHandler graphQlHandler, HttpMessageConverter<?> converter, Duration connectionInitTimeout) {
 
+		this(graphQlHandler, converter, connectionInitTimeout, null);
+	}
+
+	/**
+	 * Create a new instance.
+	 * @param graphQlHandler common handler for GraphQL over WebSocket requests
+	 * @param converter for JSON encoding and decoding
+	 * @param connectionInitTimeout how long to wait after the establishment of
+	 * the WebSocket for the {@code "connection_ini"} message from the client.
+	 * @param keepAliveDuration how frequently to send ping messages; if not
+	 * set then ping messages are not sent.
+	 * @since 1.3
+	 */
+	public GraphQlWebSocketHandler(
+			WebGraphQlHandler graphQlHandler, HttpMessageConverter<?> converter,
+			Duration connectionInitTimeout, @Nullable Duration keepAliveDuration) {
+
 		Assert.notNull(graphQlHandler, "WebGraphQlHandler is required");
 		Assert.notNull(converter, "HttpMessageConverter for JSON is required");
 
@@ -124,8 +145,10 @@ public GraphQlWebSocketHandler(
 		this.webSocketGraphQlInterceptor = this.graphQlHandler.getWebSocketInterceptor();
 		this.initTimeoutDuration = connectionInitTimeout;
 		this.converter = converter;
+		this.keepAliveDuration = keepAliveDuration;
 	}
 
+
 	@Override
 	public List<String> getSubProtocols() {
 		return SUB_PROTOCOL_LIST;
@@ -257,6 +280,21 @@ private void handleInternal(WebSocketSession session, TextMessage webSocketMessa
 							return Mono.empty();
 						})
 						.block(Duration.ofSeconds(10));
+
+				if (this.keepAliveDuration != null) {
+					Flux.interval(this.keepAliveDuration, this.keepAliveDuration)
+							.filter((aLong) -> true)
+							.publishOn(state.getScheduler()) // Serial blocking send via single thread
+							.doOnNext((aLong) -> {
+								try {
+									session.sendMessage(encode(GraphQlWebSocketMessage.ping(null)));
+								}
+								catch (IOException ex) {
+									ExceptionWebSocketHandlerDecorator.tryCloseWithError(session, ex, logger);
+								}
+							})
+							.subscribe(state.getKeepAliveSubscriber());
+				}
 			}
 			default -> GraphQlStatus.closeSession(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
 		}
@@ -444,9 +482,12 @@ private static class SessionState {
 
 		private final Scheduler scheduler;
 
-		SessionState(String graphQlSessionId, WebSocketSessionInfo sessionInfo) {
+		private final KeepAliveSubscriber keepAliveSubscriber;
+
+		SessionState(String graphQlSessionId, WebMvcSessionInfo sessionInfo) {
 			this.sessionInfo = sessionInfo;
 			this.scheduler = Schedulers.newSingle("GraphQL-WsSession-" + graphQlSessionId);
+			this.keepAliveSubscriber = new KeepAliveSubscriber(sessionInfo.getSession());
 		}
 
 		WebSocketSessionInfo getSessionInfo() {
@@ -462,12 +503,16 @@ boolean setConnectionInitPayload(Map<String, Object> payload) {
 			return this.connectionInitPayloadRef.compareAndSet(null, payload);
 		}
 
+		KeepAliveSubscriber getKeepAliveSubscriber() {
+			return this.keepAliveSubscriber;
+		}
 
 		Map<String, Subscription> getSubscriptions() {
 			return this.subscriptions;
 		}
 
 		void dispose() {
+			this.keepAliveSubscriber.cancel();
 			for (Map.Entry<String, Subscription> entry : this.subscriptions.entrySet()) {
 				try {
 					entry.getValue().cancel();
@@ -525,6 +570,10 @@ public Mono<Principal> getPrincipal() {
 		public InetSocketAddress getRemoteAddress() {
 			return this.session.getRemoteAddress();
 		}
+
+		WebSocketSession getSession() {
+			return this.session;
+		}
 	}
 
 
@@ -567,9 +616,29 @@ public void hookOnError(Throwable ex) {
 		public void hookOnComplete() {
 			this.sessionState.getSubscriptions().remove(this.subscriptionId);
 		}
+	}
+
+
+	private static class KeepAliveSubscriber extends BaseSubscriber<Long> {
+
+		private final WebSocketSession session;
+
+		KeepAliveSubscriber(WebSocketSession session) {
+			this.session = session;
+		}
 
+		@Override
+		protected void hookOnSubscribe(Subscription subscription) {
+			subscription.request(Integer.MAX_VALUE);
+		}
+
+		@Override
+		public void hookOnError(Throwable ex) {
+			ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.session, ex, logger);
+		}
 	}
 
+
 	@SuppressWarnings("serial")
 	private static final class SubscriptionExistsException extends RuntimeException {
 	}