Correct handling of complete WebSocket message

This `handleConnectionCompletion` method on `WebSocketInterceptor`
method handles the `complete` GraphQL over WebSocket message which is
defined as canceling an individual subscription stream. The method was
incorrectly modeled as a notification for the closing of the whole
connection.

The method has been renamed, Javadoc clarified, and the subscription
id passed in.

See gh-276

spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * 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.
@@ -19,6 +19,7 @@ import java.util.Map;
 
 import reactor.core.publisher.Mono;
 
+
 /**
  * An extension of {@link WebInterceptor} with additional methods to handle the
  * start and end of a WebSocket connection. Only a single interceptor of type
@@ -35,22 +36,24 @@ public interface WebSocketInterceptor extends WebInterceptor {
 	}
 
 	/**
-	 * Handle the payload from the connection initialization message that a
-	 * GraphQL over WebSocket client must send after the WebSocket session is
-	 * established and before sending any requests.
-	 * @param payload the payload from the {@code ConnectionInit} message
-	 * @return an optional payload for the {@code ConnectionAck} message
+	 * Handle the {@code "connection_init"} message at the start of a GraphQL over
+	 * WebSocket session and return an optional payload for the
+	 * {@code "connection_ack"} message to send back.
+	 * @param connectionInitPayload the payload from the {@code "connection_init"} message
+	 * @return the payload for the {@code "connection_ack"}, or empty
 	 */
-	default Mono<Object> handleConnectionInitialization(Map<String, Object> payload) {
+	default Mono<Object> handleConnectionInitialization(Map<String, Object> connectionInitPayload) {
 		return Mono.empty();
 	}
 
 	/**
-	 * Handle the completion message that a GraphQL over WebSocket clients sends
-	 * before closing the WebSocket connection.
-	 * @return signals the end of completion handling
+	 * Handle the {@code "complete"} message that clients send to stop listening
+	 * to the subscription with the given id.
+	 * <p>Note that the {@link org.reactivestreams.Publisher} for the subscription
+	 * is automatically cancelled and there is no need to do that from here.
+	 * @return {@code Mono} for the completion of handling
 	 */
-	default Mono<Void> handleConnectionCompletion() {
+	default Mono<Void> handleCancelledSubscription(String subscriptionId) {
 		return Mono.empty();
 	}
 

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

@@ -141,8 +141,9 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 						if (subscription != null) {
 							subscription.cancel();
 						}
+						return this.webSocketInterceptor.handleCancelledSubscription(id).thenMany(Flux.empty());
 					}
-					return this.webSocketInterceptor.handleConnectionCompletion().thenMany(Flux.empty());
+					return Flux.empty();
 				case "connection_init":
 					if (!connectionInitProcessed.compareAndSet(false, true)) {
 						return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);

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

@@ -169,8 +169,8 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 				if (subscription != null) {
 					subscription.cancel();
 				}
+				this.webSocketInterceptor.handleCancelledSubscription(id).block(Duration.ofSeconds(10));
 			}
-			this.webSocketInterceptor.handleConnectionCompletion().block(Duration.ofSeconds(10));
 			return;
 		case "connection_init":
 			if (sessionState.setConnectionInitProcessed()) {