Prevent starting second subscription with the same id, remove he complete message as it's not needed, and probably preventing flux from closing. (ExpediaGroup#520)

Author: gklijs

graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/execution/ApolloSubscriptionProtocolHandler.kt

@@ -65,10 +65,7 @@ class ApolloSubscriptionProtocolHandler(
                     return ackowledgeMessageFlux.concatWith(keepAliveFlux)
                 }
                 GQL_START.type -> return startSubscription(operationMessage, session)
-                GQL_STOP.type -> {
-                    sessionState.stopOperation(session, operationMessage)
-                    return Flux.empty()
-                }
+                GQL_STOP.type -> return sessionState.stopOperation(session, operationMessage)
                 GQL_CONNECTION_TERMINATE.type -> {
                     sessionState.terminateSession(session)
                     return Flux.empty()
@@ -107,6 +104,11 @@ class ApolloSubscriptionProtocolHandler(
             return Flux.just(basicConnectionErrorMessage)
         }
 
+        if (sessionState.operationExists(session, operationMessage)) {
+            logger.info("Already subscribed to operation ${operationMessage.id} for session ${session.id}")
+            return Flux.empty()
+        }
+
         val payload = operationMessage.payload
 
         if (payload == null) {
@@ -125,7 +127,6 @@ class ApolloSubscriptionProtocolHandler(
                         SubscriptionOperationMessage(type = GQL_DATA.type, id = operationMessage.id, payload = it)
                     }
                 }
-                .concatWith(Flux.just(SubscriptionOperationMessage(type = SubscriptionOperationMessage.ServerMessages.GQL_COMPLETE.type, id = operationMessage.id)))
                 .doOnSubscribe { sessionState.saveOperation(session, operationMessage, it) }
         } catch (exception: Exception) {
             logger.error("Error running graphql subscription", exception)

graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/execution/ApolloSubscriptionSessionState.kt

@@ -17,8 +17,10 @@
 package com.expediagroup.graphql.spring.execution
 
 import com.expediagroup.graphql.spring.model.SubscriptionOperationMessage
+import com.expediagroup.graphql.spring.model.SubscriptionOperationMessage.ServerMessages.GQL_COMPLETE
 import org.reactivestreams.Subscription
 import org.springframework.web.reactive.socket.WebSocketSession
+import reactor.core.publisher.Flux
 import java.util.concurrent.ConcurrentHashMap
 
 internal class ApolloSubscriptionSessionState {
@@ -53,16 +55,19 @@ internal class ApolloSubscriptionSessionState {
     /**
      * Stop the subscription sending data. Does NOT terminate the session.
      */
-    fun stopOperation(session: WebSocketSession, operationMessage: SubscriptionOperationMessage) {
+    fun stopOperation(session: WebSocketSession, operationMessage: SubscriptionOperationMessage): Flux<SubscriptionOperationMessage> {
         if (operationMessage.id != null) {
             val operationsForSession = activeOperations[session.id]
-            operationsForSession?.get(operationMessage.id)?.cancel()
-            operationsForSession?.remove(operationMessage.id)
-
-            if (operationsForSession?.isEmpty() == true) {
-                activeOperations.remove(session.id)
+            operationsForSession?.get(operationMessage.id)?.let {
+                it.cancel()
+                operationsForSession.remove(operationMessage.id)
+                if (operationsForSession.isEmpty()) {
+                    activeOperations.remove(session.id)
+                }
+                return Flux.just(SubscriptionOperationMessage(type = GQL_COMPLETE.type, id = operationMessage.id))
             }
         }
+        return Flux.empty()
     }
 
     /**
@@ -75,4 +80,10 @@ internal class ApolloSubscriptionSessionState {
         activeKeepAliveSessions.remove(session.id)
         session.close()
     }
+
+    /**
+     * Looks up the operation for the client, to check if it already exists
+     */
+    fun operationExists(session: WebSocketSession, operationMessage: SubscriptionOperationMessage): Boolean =
+        activeOperations[session.id]?.containsKey(operationMessage.id) ?: false
 }

graphql-kotlin-spring-server/src/test/kotlin/com/expediagroup/graphql/spring/execution/ApolloSubscriptionProtocolHandlerTest.kt

@@ -244,11 +244,11 @@ class ApolloSubscriptionProtocolHandlerTest {
     fun `Return GQL_CONNECTION_ERROR when sending GQL_START but id is null`() {
         val config: GraphQLConfigurationProperties = mockk()
         val operationMessage = SubscriptionOperationMessage(type = GQL_START.type, id = null).toJson()
-        val session: WebSocketSession = mockk()
+        val mockSession: WebSocketSession = mockk { every { id } returns "123" }
         val subscriptionHandler: SubscriptionHandler = mockk()
 
         val handler = ApolloSubscriptionProtocolHandler(config, subscriptionHandler, objectMapper)
-        val flux = handler.handle(operationMessage, session)
+        val flux = handler.handle(operationMessage, mockSession)
 
         val message = flux.blockFirst(Duration.ofSeconds(2))
         assertNotNull(message)
@@ -316,7 +316,60 @@ class ApolloSubscriptionProtocolHandlerTest {
         val graphQLResponse: GraphQLResponse = objectMapper.convertValue(payload)
         assertEquals(expected = "myData", actual = graphQLResponse.data)
 
-        assertEquals(expected = 2, actual = flux.count().block())
+        assertEquals(expected = 1, actual = flux.count().block())
+        verify(exactly = 0) { session.close() }
+    }
+
+    @Test
+    fun `Return GQL_COMPLETE when sending GQL_STOP with GraphQLRequest having operation id of running operation`() {
+        val config: GraphQLConfigurationProperties = mockk()
+        val graphQLRequest = GraphQLRequest("{ message }")
+        val startRequest = SubscriptionOperationMessage(type = GQL_START.type, id = "abc", payload = graphQLRequest).toJson()
+        val stopRequest = SubscriptionOperationMessage(type = GQL_STOP.type, id = "abc").toJson()
+        val session: WebSocketSession = mockk {
+            every { close() } returns mockk()
+            every { id } returns "123"
+        }
+        val subscriptionHandler: SubscriptionHandler = mockk {
+            every { executeSubscription(eq(graphQLRequest)) } returns Flux.just(GraphQLResponse("myData"))
+        }
+
+        val handler = ApolloSubscriptionProtocolHandler(config, subscriptionHandler, objectMapper)
+        val startFlux = handler.handle(startRequest, session)
+        startFlux.blockFirst(Duration.ofSeconds(2))
+        val stopFlux = handler.handle(stopRequest, session)
+
+        StepVerifier.create(stopFlux)
+            .expectSubscription()
+            .expectNextMatches { it.type == "complete" }
+            .thenCancel()
+            .verify()
+
+        assertEquals(expected = 1, actual = startFlux.count().block())
+        assertEquals(expected = 1, actual = stopFlux.count().block())
+        verify(exactly = 0) { session.close() }
+    }
+
+    @Test
+    fun `Dont start second subscription when operation id is already in activeOperations`() {
+        val config: GraphQLConfigurationProperties = mockk()
+        val graphQLRequest = GraphQLRequest("{ message }")
+        val operationMessage = SubscriptionOperationMessage(type = SubscriptionOperationMessage.ClientMessages.GQL_START.type, id = "abc", payload = graphQLRequest).toJson()
+        val session: WebSocketSession = mockk {
+            every { close() } returns mockk()
+            every { id } returns "123"
+        }
+        val subscriptionHandler: SubscriptionHandler = mockk {
+            every { executeSubscription(eq(graphQLRequest)) } returns Flux.just(GraphQLResponse("myData"))
+        }
+
+        val handler = ApolloSubscriptionProtocolHandler(config, subscriptionHandler, objectMapper)
+        val flux = handler.handle(operationMessage, session)
+        flux.blockFirst(Duration.ofSeconds(2))
+        val fluxTwo = handler.handle(operationMessage, session)
+
+        assertEquals(expected = 1, actual = flux.count().block())
+        assertEquals(expected = 0, actual = fluxTwo.count().block())
         verify(exactly = 0) { session.close() }
     }
 
@@ -337,7 +390,7 @@ class ApolloSubscriptionProtocolHandlerTest {
         val handler = ApolloSubscriptionProtocolHandler(config, subscriptionHandler, objectMapper)
         val flux = handler.handle(operationMessage, session)
 
-        assertEquals(expected = 2, actual = flux.count().block())
+        assertEquals(expected = 1, actual = flux.count().block())
         val message = flux.blockFirst(Duration.ofSeconds(2))
         assertNotNull(message)
         assertEquals(expected = GQL_ERROR.type, actual = message.type)