recover from panic

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'")
 	}
-	dataChan_ <- wsResp
+	graphql.WriteToChannelOrRecover(dataChan_, wsResp)
 	return nil
 }
 {{end}}

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

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

graphql/client.go

@@ -131,6 +131,26 @@ 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)
@@ -140,22 +160,13 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if wsMsg.ID == "" { // e.g. keep-alive messages
 		return nil
 	}
-	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
+	sub, err := w.getSubscriptionOrHandleComplete(wsMsg.ID, wsMsg.Type)
+	if err != nil {
+		return err
 	}
-	if wsMsg.Type == webSocketTypeComplete {
-		sub.hasBeenUnsubscribed = true
-		w.subscriptions.map_[wsMsg.ID] = sub
-		reflect.ValueOf(sub.interfaceChan).Close()
+	if sub == nil {
 		return nil
 	}
-
 	return sub.forwardDataFunc(sub.interfaceChan, wsMsg.Payload)
 }
 

graphql/websocket.go

@@ -164,7 +164,6 @@ func countTo(ctx context.Context, stopCount int) (<-chan int, error) {
 		for counter := range stopCount {
 			select {
 			case <-ctx.Done():
-				fmt.Println("ctx done:", ctx.Err())
 				return
 			default:
 				respChan <- counter