websocket: prevent deadlocks

.github/workflows/go.yml

@@ -29,7 +29,7 @@ jobs:
         # Needed for the example-test to run.
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: |
-        go test -cover -v ./...
+        go test -race -cover -v ./...
 
   lint:
     name: Lint

Makefile

@@ -7,7 +7,7 @@ lint:
 	internal/lint/golangci-lint run ./... --fix
 
 check: lint
-	go test -cover ./...
+	go test -race -cover ./...
 	go mod tidy
 
 .PHONY: example

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.SafeSend(dataChan_, wsResp)
 	return nil
 }
 {{end}}

graphql/client.go

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

graphql/subscription.go

@@ -37,13 +37,8 @@ func (s *subscriptionMap) Unsubscribe(subscriptionID string) error {
 	if !success {
 		return fmt.Errorf("tried to unsubscribe from unknown subscription with ID '%s'", subscriptionID)
 	}
-	hasBeenUnsubscribed := unsub.hasBeenUnsubscribed
 	unsub.hasBeenUnsubscribed = true
 	s.map_[subscriptionID] = unsub
-
-	if !hasBeenUnsubscribed {
-		reflect.ValueOf(s.map_[subscriptionID].interfaceChan).Close()
-	}
 	return nil
 }
 
@@ -61,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/websocket.go

@@ -6,9 +6,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"reflect"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -51,7 +51,7 @@ type webSocketClient struct {
 	connParams    map[string]interface{}
 	errChan       chan error
 	subscriptions subscriptionMap
-	isClosing     bool
+	isClosing     atomic.Bool
 	sync.Mutex
 }
 
@@ -107,14 +107,14 @@ func (w *webSocketClient) waitForConnAck() error {
 func (w *webSocketClient) handleErr(err error) {
 	w.Lock()
 	defer w.Unlock()
-	if !w.isClosing {
+	if !w.isClosing.Load() {
 		w.errChan <- err
 	}
 }
 
 func (w *webSocketClient) listenWebSocket() {
 	for {
-		if w.isClosing {
+		if w.isClosing.Load() {
 			return
 		}
 		_, message, err := w.conn.ReadMessage()
@@ -139,22 +139,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.subscriptions.GetOrClose(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)
 }
 
@@ -208,7 +199,7 @@ func (w *webSocketClient) Close() error {
 	}
 	w.Lock()
 	defer w.Unlock()
-	w.isClosing = true
+	w.isClosing.Store(true)
 	close(w.errChan)
 	return w.conn.Close()
 }

internal/integration/integration_test.go

@@ -252,10 +252,6 @@ func TestSubscriptionClose(t *testing.T) {
 	_ = `# @genqlient
 	subscription countClose { countClose }`
 
-	ctx := context.Background()
-	server := server.RunServer()
-	defer server.Close()
-
 	cases := []struct {
 		name  string
 		unsub bool
@@ -272,6 +268,12 @@ func TestSubscriptionClose(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
+			server := server.RunServer()
+			defer server.Close()
+
 			wsClient := newRoundtripWebSocketClient(t, server.URL)
 
 			_, err := wsClient.Start(ctx)

internal/integration/server/server.go

@@ -157,33 +157,37 @@ func (m mutationResolver) CreateUser(ctx context.Context, input NewUser) (*User,
 	return &newUser, nil
 }
 
-func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
+func countTo(ctx context.Context, stopCount int) (<-chan int, error) {
 	respChan := make(chan int, 1)
 	go func(respChan chan int) {
 		defer close(respChan)
-		counter := 0
-		for {
-			if counter == 10 {
+		for counter := range stopCount {
+			select {
+			case <-ctx.Done():
 				return
+			default:
+				respChan <- counter
+				time.Sleep(100 * time.Millisecond)
 			}
-			respChan <- counter
-			counter++
-			time.Sleep(100 * time.Millisecond)
 		}
 	}(respChan)
 	return respChan, nil
 }
 
+func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
+	return countTo(ctx, 10)
+}
+
 func (s *subscriptionResolver) CountAuthorized(ctx context.Context) (<-chan int, error) {
 	if getAuthToken(ctx) != "authorized-user-token" {
 		return nil, fmt.Errorf("unauthorized")
 	}
 
-	return s.Count(ctx)
+	return countTo(ctx, 10)
 }
 
 func (s *subscriptionResolver) CountClose(ctx context.Context) (<-chan int, error) {
-	return s.Count(ctx)
+	return countTo(ctx, 1000)
 }
 
 const AuthKey = "authToken"