Tie up loose ends with subscription handling

Author: apottere

src/main/java/graphql/servlet/internal/ApolloSubscriptionProtocolHandler.java

@@ -7,17 +7,21 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.websocket.CloseReason;
 import javax.websocket.Session;
 import javax.websocket.server.HandshakeRequest;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 
 import static graphql.servlet.internal.ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_COMPLETE;
+import static graphql.servlet.internal.ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_CONNECTION_TERMINATE;
 import static graphql.servlet.internal.ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_DATA;
 import static graphql.servlet.internal.ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_ERROR;
 
 /**
+ * https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
+ *
  * @author Andrew Potter
  */
 public class ApolloSubscriptionProtocolHandler extends SubscriptionProtocolHandler {
@@ -57,6 +61,21 @@ public void onMessage(HandshakeRequest request, Session session, WsSessionSubscr
                     ))
                 );
                 break;
+
+            case GQL_STOP:
+                unsubscribe(subscriptions, message.id);
+                break;
+
+            case GQL_CONNECTION_TERMINATE:
+                try {
+                    session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "client requested " + GQL_CONNECTION_TERMINATE.getType()));
+                } catch (IOException e) {
+                    log.error("Unable to close websocket session!", e);
+                }
+                break;
+
+            default:
+                throw new IllegalArgumentException("Unknown message type: " + message.getType());
         }
     }
 

src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolHandler.java

@@ -3,6 +3,7 @@
 import javax.websocket.Session;
 import javax.websocket.server.HandshakeRequest;
 import java.io.IOException;
+import java.util.UUID;
 
 /**
  * @author Andrew Potter
@@ -17,8 +18,16 @@ public FallbackSubscriptionProtocolHandler(SubscriptionHandlerInput subscription
 
     @Override
     public void onMessage(HandshakeRequest request, Session session, WsSessionSubscriptions subscriptions, String text) throws Exception {
-        subscribe(session, input.getQueryInvoker().query(input.getInvocationInputFactory().create(
-                input.getGraphQLObjectMapper().readGraphQLRequest(text))), subscriptions, session.getId());
+        subscribe(
+            session,
+            input.getQueryInvoker().query(
+                input.getInvocationInputFactory().create(
+                    input.getGraphQLObjectMapper().readGraphQLRequest(text)
+                )
+            ),
+            subscriptions,
+            UUID.randomUUID().toString()
+        );
     }
 
     @Override

src/main/java/graphql/servlet/internal/SubscriptionProtocolHandler.java

@@ -33,7 +33,7 @@ protected void subscribe(Session session, ExecutionResult executionResult, WsSes
 
         if (data instanceof Publisher) {
             @SuppressWarnings("unchecked") final Publisher<ExecutionResult> publisher = (Publisher<ExecutionResult>) data;
-            final AtomicReference<Subscription> subscriptionReference = new AtomicReference<>();
+            final AtomicSubscriptionReference subscriptionReference = new AtomicSubscriptionReference();
 
             publisher.subscribe(new Subscriber<ExecutionResult>() {
                 @Override
@@ -57,16 +57,39 @@ public void onError(Throwable throwable) {
                     log.error("Subscription error", throwable);
                     subscriptions.cancel(id);
                     sendErrorMessage(session, id);
-//                    sendMessage(session, ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_ERROR, id);
                 }
 
                 @Override
                 public void onComplete() {
                     subscriptions.cancel(id);
                     sendCompleteMessage(session, id);
-//                    sendMessage(session, ApolloSubscriptionProtocolHandler.OperationMessage.Type.GQL_COMPLETE, id);
                 }
             });
         }
     }
+
+    protected void unsubscribe(WsSessionSubscriptions subscriptions, String id) {
+        subscriptions.cancel(id);
+    }
+
+    static class AtomicSubscriptionReference {
+        private final AtomicReference<Subscription> reference = new AtomicReference<>(null);
+
+        public void set(Subscription subscription) {
+            if(reference.get() != null) {
+                throw new IllegalStateException("Cannot overwrite subscription!");
+            }
+
+            reference.set(subscription);
+        }
+
+        public Subscription get() {
+            Subscription subscription = reference.get();
+            if(subscription == null) {
+                throw new IllegalStateException("Subscription has not been initialized yet!");
+            }
+
+            return subscription;
+        }
+    }
 }