websocket: gracefully close connection

graphql/subscription.go

@@ -30,13 +30,6 @@ func (s *subscriptionMap) Create(subscriptionID string, interfaceChan interface{
 	}
 }
 
-func (s *subscriptionMap) Read(subscriptionID string) (sub subscription, success bool) {
-	s.RLock()
-	defer s.RUnlock()
-	sub, success = s.map_[subscriptionID]
-	return sub, success
-}
-
 func (s *subscriptionMap) Unsubscribe(subscriptionID string) error {
 	s.Lock()
 	defer s.Unlock()

graphql/websocket.go

@@ -139,14 +139,18 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if wsMsg.ID == "" { // e.g. keep-alive messages
 		return nil
 	}
-	sub, ok := w.subscriptions.Read(wsMsg.ID)
-	if !ok {
+	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
 	}
 	if wsMsg.Type == webSocketTypeComplete {
+		sub.hasBeenUnsubscribed = true
+		w.subscriptions.map_[wsMsg.ID] = sub
 		reflect.ValueOf(sub.interfaceChan).Close()
 		return nil
 	}
@@ -194,13 +198,13 @@ func (w *webSocketClient) Close() error {
 	if w.conn == nil {
 		return nil
 	}
-	err := w.conn.WriteMessage(closeMessage, formatCloseMessage(closeNormalClosure, ""))
+	err := w.UnsubscribeAll()
 	if err != nil {
-		return fmt.Errorf("failed to send closure message: %w", err)
+		return fmt.Errorf("failed to unsubscribe: %w", err)
 	}
-	err = w.UnsubscribeAll()
+	err = w.conn.WriteMessage(closeMessage, formatCloseMessage(closeNormalClosure, ""))
 	if err != nil {
-		return fmt.Errorf("failed to unsubscribe: %w", err)
+		return fmt.Errorf("failed to send closure message: %w", err)
 	}
 	w.Lock()
 	defer w.Unlock()

internal/integration/integration_test.go

@@ -104,7 +104,10 @@ func TestSubscription(t *testing.T) {
 
 			dataChan, subscriptionID, err := count(ctx, wsClient)
 			require.NoError(t, err)
-			defer wsClient.Close()
+			defer func() {
+				err := wsClient.Close()
+				require.NoError(t, err)
+			}()
 
 			var (
 				counter = 0
@@ -198,7 +201,10 @@ func TestSubscriptionConnectionParams(t *testing.T) {
 
 			dataChan, subscriptionID, err := countAuthorized(ctx, wsClient)
 			require.NoError(t, err)
-			defer wsClient.Close()
+			defer func() {
+				err := wsClient.Close()
+				require.NoError(t, err)
+			}()
 
 			var (
 				counter = 0
@@ -242,6 +248,49 @@ func TestSubscriptionConnectionParams(t *testing.T) {
 	}
 }
 
+func TestSubscriptionClose(t *testing.T) {
+	_ = `# @genqlient
+	subscription countClose { countClose }`
+
+	ctx := context.Background()
+	server := server.RunServer()
+	defer server.Close()
+
+	cases := []struct {
+		name  string
+		unsub bool
+	}{
+		{
+			name:  "unsubscribed_manually",
+			unsub: true,
+		},
+		{
+			name:  "unsubscribed_automatically",
+			unsub: false,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			wsClient := newRoundtripWebSocketClient(t, server.URL)
+
+			_, err := wsClient.Start(ctx)
+			require.NoError(t, err)
+
+			_, subscriptionID, err := countClose(ctx, wsClient)
+			require.NoError(t, err)
+
+			if tc.unsub {
+				err = wsClient.Unsubscribe(subscriptionID)
+				require.NoError(t, err)
+			}
+
+			err = wsClient.Close()
+			require.NoError(t, err)
+		})
+	}
+}
+
 func TestServerError(t *testing.T) {
 	_ = `# @genqlient
 	query failingQuery { fail me { id } }`

internal/integration/schema.graphql

@@ -20,6 +20,7 @@ type Mutation {
 type Subscription {
   count: Int!
   countAuthorized: Int!
+  countClose: Int!
 }
 
 type User implements Being & Lucky {

internal/integration/server/server.go

@@ -182,6 +182,10 @@ func (s *subscriptionResolver) CountAuthorized(ctx context.Context) (<-chan int,
 	return s.Count(ctx)
 }
 
+func (s *subscriptionResolver) CountClose(ctx context.Context) (<-chan int, error) {
+	return s.Count(ctx)
+}
+
 const AuthKey = "authToken"
 
 type (