handle concurrent unsubs

Author: HaraldNordgren

graphql/subscription.go

@@ -30,22 +30,6 @@ func (s *subscriptionMap) Create(subscriptionID string, interfaceChan interface{
 	}
 }
 
-func (s *subscriptionMap) Read(subscriptionID string) (sub subscription, success bool) {
-	s.RLock()
-	defer s.RUnlock()
-	sub, success = s.map_[subscriptionID]
-	return sub, success
-}
-
-func (s *subscriptionMap) markUnsubscribed(subscriptionID string) {
-	s.Lock()
-	defer s.Unlock()
-	if sub, ok := s.map_[subscriptionID]; ok {
-		sub.hasBeenUnsubscribed = true
-		s.map_[subscriptionID] = sub
-	}
-}
-
 func (s *subscriptionMap) Unsubscribe(subscriptionID string) error {
 	s.Lock()
 	defer s.Unlock()

graphql/websocket.go

@@ -139,15 +139,18 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if wsMsg.ID == "" { // e.g. keep-alive messages
 		return nil
 	}
-	sub, ok := w.subscriptions.Read(wsMsg.ID)
-	if !ok {
+	w.subscriptions.Lock()
+	defer w.subscriptions.Unlock()
+	sub, success := w.subscriptions.map_[wsMsg.ID]
+	if !success {
 		return fmt.Errorf("received message for unknown subscription ID '%s'", wsMsg.ID)
 	}
 	if sub.hasBeenUnsubscribed {
 		return nil
 	}
 	if wsMsg.Type == webSocketTypeComplete {
-		w.subscriptions.markUnsubscribed(wsMsg.ID)
+		sub.hasBeenUnsubscribed = true
+		w.subscriptions.map_[wsMsg.ID] = sub
 		reflect.ValueOf(sub.interfaceChan).Close()
 		return nil
 	}

internal/integration/integration_test.go

@@ -104,7 +104,9 @@ func TestSubscription(t *testing.T) {
 
 			dataChan, subscriptionID, err := count(ctx, wsClient)
 			require.NoError(t, err)
-			defer wsClient.Close()
+			defer func() {
+				require.NoError(t, wsClient.Close())
+			}()
 
 			var (
 				counter = 0
@@ -198,7 +200,9 @@ func TestSubscriptionConnectionParams(t *testing.T) {
 
 			dataChan, subscriptionID, err := countAuthorized(ctx, wsClient)
 			require.NoError(t, err)
-			defer wsClient.Close()
+			defer func() {
+				require.NoError(t, wsClient.Close())
+			}()
 
 			var (
 				counter = 0