refactor

Author: HaraldNordgren

generate/operation.go.tmpl

@@ -74,7 +74,7 @@ func {{.Name}}ForwardData(interfaceChan interface{}, jsonRawMsg json.RawMessage)
 	if !ok {
 		return errors.New("failed to cast interface into 'chan {{.Name}}WsResponse'")
 	}
-	graphql.WriteToChannelOrRecover(dataChan_, wsResp)
+	graphql.SafeSend(dataChan_, wsResp)
 	return nil
 }
 {{end}}

generate/testdata/snapshots/TestGenerate-SimpleSubscription.graphql-SimpleSubscription.graphql.go

@@ -364,7 +364,7 @@ func (c *client) createGetRequest(req *Request) (*http.Request, error) {
 	return httpReq, nil
 }
 
-func WriteToChannelOrRecover[T any](dataChan_ chan T, wsResp T) {
+func SafeSend[T any](dataChan_ chan T, wsResp T) {
 	defer func() {
 		_ = recover()
 	}()

graphql/client.go

@@ -2,6 +2,7 @@ package graphql
 
 import (
 	"fmt"
+	"reflect"
 	"sync"
 )
 
@@ -55,3 +56,23 @@ func (s *subscriptionMap) Delete(subscriptionID string) {
 	defer s.Unlock()
 	delete(s.map_, subscriptionID)
 }
+
+func (s *subscriptionMap) GetOrClose(subscriptionID string, subscriptionType string) (*subscription, error) {
+	s.Lock()
+	defer s.Unlock()
+	sub, success := s.map_[subscriptionID]
+	if !success {
+		return nil, fmt.Errorf("received message for unknown subscription ID '%s'", subscriptionID)
+	}
+	if sub.hasBeenUnsubscribed {
+		return nil, nil
+	}
+	if subscriptionType == webSocketTypeComplete {
+		sub.hasBeenUnsubscribed = true
+		s.map_[subscriptionID] = sub
+		reflect.ValueOf(sub.interfaceChan).Close()
+		return nil, nil
+	}
+
+	return &sub, nil
+}

graphql/subscription.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"reflect"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -131,26 +130,6 @@ func (w *webSocketClient) listenWebSocket() {
 	}
 }
 
-func (w *webSocketClient) getSubscriptionOrHandleComplete(subscriptionID string, subscriptionType string) (*subscription, error) {
-	w.subscriptions.Lock()
-	defer w.subscriptions.Unlock()
-	sub, success := w.subscriptions.map_[subscriptionID]
-	if !success {
-		return nil, fmt.Errorf("received message for unknown subscription ID '%s'", subscriptionID)
-	}
-	if sub.hasBeenUnsubscribed {
-		return nil, nil
-	}
-	if subscriptionType == webSocketTypeComplete {
-		sub.hasBeenUnsubscribed = true
-		w.subscriptions.map_[subscriptionID] = 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)
@@ -160,7 +139,7 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if wsMsg.ID == "" { // e.g. keep-alive messages
 		return nil
 	}
-	sub, err := w.getSubscriptionOrHandleComplete(wsMsg.ID, wsMsg.Type)
+	sub, err := w.subscriptions.GetOrClose(wsMsg.ID, wsMsg.Type)
 	if err != nil {
 		return err
 	}