Fix deadloack introduced in #402

This PR fixes the deadlock introduced in #402. The deadlock was occurring
because the forwardWebSocketData function was holding a lock on the
subscription map (in the listenWebSocket goroutine) while Unsubscribe was
blocked attempting to aquire the same lock. The forwardWebSocketData function
was also blocked attempting to send data on the interfaceChan channel, which
had no readers in the test being run.

After fixing the locking (by making sure the lock isn't held when
forwardDataFunc is called), a race condition cropped up between sending on
interfaceChan (in forwardWebSocketData) and closing interfaceChan in
Unsubscribe. This was fixed by making the listenWebSocket goroutine the "owner"
of the interfaceChan channel -- it is now the only goroutine that sends on and
closes the channel. Other goroutines singal interfaceChan should be closed by
setting `_hasBeenUnsubscribed` on the channel, which is protected by a
short-held lock.

Note that there's a similar data race involving isClosing on webSocketClient.
isClosing is set when Close is called (while holding a lock) but read in
listenWebSocket without a lock. A separate PR will fix this race (so the
`-race` flag can be added when running tests).

Author: dnerdy

graphql/subscription.go

@@ -2,7 +2,6 @@ package graphql
 
 import (
 	"fmt"
-	"reflect"
 	"sync"
 )
 
@@ -13,10 +12,40 @@ type subscriptionMap struct {
 }
 
 type subscription struct {
+	// interfaceChan is passed in by the user when creating a subscription but
+	// closed by webSocketClient when the subscription is unsubscribed, i.e.
+	// ownership of interfaceChan is passed from the user to the client.
+	//
+	// The subscription is unsubscribed either explicitly by the user or when
+	// a message of webSocketTypeComplete is received. On unsubscribe,
+	// the _hasBeenUnsubscribed flag is set to true. listenWebSocket then
+	// closes interfaceChan on the next receive loop.
+	//
+	// The listenWebSocket client method handles both sending on the channel
+	// and closing of the channel, so is no possibility of races between send
+	// and close.
 	interfaceChan       interface{}
+
 	forwardDataFunc     ForwardDataFunction
 	id                  string
-	hasBeenUnsubscribed bool
+
+	// Hold when accessing _hasBeenUnsubscribed
+	hasBeenUnsubscribedMu sync.Mutex
+	_hasBeenUnsubscribed bool
+}
+
+func (s *subscription) unsubscribe() {
+	s.hasBeenUnsubscribedMu.Lock()
+	defer s.hasBeenUnsubscribedMu.Unlock()
+
+	s._hasBeenUnsubscribed = true
+}
+
+func (s *subscription) hasBeenUnsubscribed() bool {
+	s.hasBeenUnsubscribedMu.Lock()
+	defer s.hasBeenUnsubscribedMu.Unlock()
+
+	return s._hasBeenUnsubscribed
 }
 
 func (s *subscriptionMap) Create(subscriptionID string, interfaceChan interface{}, forwardDataFunc ForwardDataFunction) {
@@ -26,7 +55,7 @@ func (s *subscriptionMap) Create(subscriptionID string, interfaceChan interface{
 		id:                  subscriptionID,
 		interfaceChan:       interfaceChan,
 		forwardDataFunc:     forwardDataFunc,
-		hasBeenUnsubscribed: false,
+		_hasBeenUnsubscribed: false,
 	}
 }
 
@@ -37,16 +66,28 @@ 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
+	unsub.unsubscribe()
 	s.map_[subscriptionID] = unsub
 
-	if !hasBeenUnsubscribed {
-		reflect.ValueOf(s.map_[subscriptionID].interfaceChan).Close()
-	}
 	return nil
 }
 
+func (s *subscriptionMap) forEachSubscription(fn func(sub subscription)) {
+	s.Lock()
+	defer s.Unlock()
+
+	for id := range s.map_ {
+		fn(s.map_[id])
+	}
+}
+
+func (s *subscriptionMap) GetSubscription(subscriptionID string) (subscription, bool) {
+	s.Lock()
+	defer s.Unlock()
+	sub, ok := s.map_[subscriptionID]
+	return sub, ok
+}
+
 func (s *subscriptionMap) GetAllIDs() (subscriptionIDs []string) {
 	s.RLock()
 	defer s.RUnlock()

graphql/subscription_test.go

@@ -22,7 +22,7 @@ func Test_subscriptionMap_Unsubscribe(t *testing.T) {
 						id:                  "sub1",
 						interfaceChan:       make(chan struct{}),
 						forwardDataFunc:     nil,
-						hasBeenUnsubscribed: false,
+						_hasBeenUnsubscribed: false,
 					},
 				},
 			},
@@ -45,7 +45,7 @@ func Test_subscriptionMap_Unsubscribe(t *testing.T) {
 						id:                  "sub2",
 						interfaceChan:       nil,
 						forwardDataFunc:     nil,
-						hasBeenUnsubscribed: true,
+						_hasBeenUnsubscribed: true,
 					},
 				},
 			},

graphql/websocket.go

@@ -114,6 +114,17 @@ func (w *webSocketClient) handleErr(err error) {
 
 func (w *webSocketClient) listenWebSocket() {
 	for {
+		// The listenWebSocket goroutine "owns" interfaceChan. Both sending
+		// (in forwardWebSocketData below) and closure (here) happen in this
+		// goroutine, so there is no possibility of races between send and close.
+		//
+		// interfaceChan's are closed at the top of listenWebSocket to
+		// guarantee the channels are closed even if listenWebSocket will exit.
+		w.subscriptions.forEachSubscription(func(sub subscription) {
+			if sub.hasBeenUnsubscribed() {
+				reflect.ValueOf(sub.interfaceChan).Close()
+			}
+		})
 		if w.isClosing {
 			return
 		}
@@ -139,22 +150,20 @@ 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 wsMsg.Type == webSocketTypeComplete {
+		return w.subscriptions.Unsubscribe(wsMsg.ID)
 	}
-	if sub.hasBeenUnsubscribed {
-		return nil
+
+	sub, ok := w.subscriptions.GetSubscription(wsMsg.ID)
+	if !ok {
+		return fmt.Errorf("received message for unknown subscription ID '%s'", sub.id)
 	}
-	if wsMsg.Type == webSocketTypeComplete {
-		sub.hasBeenUnsubscribed = true
-		w.subscriptions.map_[wsMsg.ID] = sub
-		reflect.ValueOf(sub.interfaceChan).Close()
+	// Note: there's no data race between hasBeenUnsubscribed and the closed
+	// state of interfaceChan because interfaceChan is only closed by the
+	// caller of this function.
+	if sub.hasBeenUnsubscribed() {
 		return nil
 	}
-
 	return sub.forwardDataFunc(sub.interfaceChan, wsMsg.Payload)
 }
 

graphql/websocket_test.go

@@ -14,7 +14,7 @@ func forgeTestWebSocketClient(hasBeenUnsubscribed bool) *webSocketClient {
 			RWMutex: sync.RWMutex{},
 			map_: map[string]subscription{
 				testSubscriptionID: {
-					hasBeenUnsubscribed: hasBeenUnsubscribed,
+					_hasBeenUnsubscribed: hasBeenUnsubscribed,
 					interfaceChan:       make(chan any),
 					forwardDataFunc: func(interfaceChan any, jsonRawMsg json.RawMessage) error {
 						return nil