move into map

Author: HaraldNordgren

graphql/subscription.go

@@ -61,3 +61,23 @@ func (s *subscriptionMap) Delete(subscriptionID string) {
 	defer s.Unlock()
 	delete(s.map_, subscriptionID)
 }
+
+func (s *subscriptionMap) GetOrHandleCompletion(id string, subType string) (*subscription, error) {
+	s.Lock()
+	defer s.Unlock()
+	sub, ok := s.map_[id]
+	if !ok {
+		return nil, fmt.Errorf("received message for unknown subscription ID '%s'", id)
+	}
+	if sub.hasBeenUnsubscribed {
+		return nil, nil
+	}
+	if subType == webSocketTypeComplete {
+		sub.hasBeenUnsubscribed = true
+		s.map_[id] = sub
+		reflect.ValueOf(sub.interfaceChan).Close()
+		return nil, nil
+	}
+
+	return &sub, nil
+}

graphql/websocket.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -130,26 +129,6 @@ func (w *webSocketClient) listenWebSocket() {
 	}
 }
 
-func (w *webSocketClient) getSubAndHandleCompletion(wsMsg webSocketReceiveMessage) (*subscription, error) {
-	w.subscriptions.Lock()
-	defer w.subscriptions.Unlock()
-	sub, success := w.subscriptions.map_[wsMsg.ID]
-	if !success {
-		return nil, fmt.Errorf("received message for unknown subscription ID '%s'", wsMsg.ID)
-	}
-	if sub.hasBeenUnsubscribed {
-		return nil, nil
-	}
-	if wsMsg.Type == webSocketTypeComplete {
-		sub.hasBeenUnsubscribed = true
-		w.subscriptions.map_[wsMsg.ID] = sub
-		reflect.ValueOf(sub.interfaceChan).Close()
-		return nil, nil
-	}
-
-	return &sub, nil
-}
-
 func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	var wsMsg webSocketReceiveMessage
 	err := json.Unmarshal(message, &wsMsg)
@@ -159,14 +138,10 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if wsMsg.ID == "" { // e.g. keep-alive messages
 		return nil
 	}
-
-	sub, err := w.getSubAndHandleCompletion(wsMsg)
-	if err != nil {
+	sub, err := w.subscriptions.GetOrHandleCompletion(wsMsg.ID, wsMsg.Type)
+	if err != nil || sub == nil {
 		return err
 	}
-	if sub == nil {
-		return nil
-	}
 
 	return sub.forwardDataFunc(sub.interfaceChan, wsMsg.Payload)
 }