Repository: Khan/genqlient
Pull request #354 (2024-09-27 10:23:37)

Fix issue where subscription server data channel closing gives error
<!--
Thanks for your contribution! Check out the
[contributing docs](https://github.com/Khan/genqlient/blob/main/docs/CONTRIBUTING.md)
for more on contributing to genqlient.
-->

Handles errors discussed here: https://github.com/Khan/genqlient/pull/347#issuecomment-2376667291


I have:
- [x] Written a clear PR title and description (above)
- [x] Signed the [Khan Academy CLA](https://www.khanacademy.org/r/cla)
- [x] Added tests covering my changes, if applicable
- [x] Included a link to the issue fixed, if applicable
- [x] Included documentation, for new features
- [x] Added an entry to the changelog


HaraldNordgren: Yeah, I don't love the reflect part, although it does solves the problem. The reason I still went with it was because I found the same thing in subscriptionsMap.Unsubscribe().

If we really don't want to reflect then a type cast could likely also solve the problem 🤗
benjaminjkraft: Ah, I guess it's not the first then!

A cast would be better, but won't work because we don't know the concrete type, right?
HaraldNordgren: Ah, right!
matthieu4294967296moineau: > Ah, I guess it's not the first then!
> 
> A cast would be better, but won't work because we don't know the concrete type, right?

yes this is why I used it in a first place to close the channel when unsubscribing

The fix looks good to me, I could test it and it worked fine 👍 
HaraldNordgren: @benjaminjkraft @matthieu4294967296moineau  Great! Thanks for the feedback!

Let's merge? 🚀

Patch:
diff --git a/graphql/websocket.go b/graphql/websocket.go
index bb9a98b..10b77da 100644
--- a/graphql/websocket.go
+++ b/graphql/websocket.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -135,6 +136,11 @@ func (w *webSocketClient) forwardWebSocketData(message []byte) error {
 	if sub.hasBeenUnsubscribed {
 		return nil
 	}
+	if wsMsg.Type == webSocketTypeComplete {
+		reflect.ValueOf(sub.interfaceChan).Close()
+		return nil
+	}
+
 	return sub.forwardDataFunc(sub.interfaceChan, wsMsg.Payload)
 }
 
diff --git a/internal/integration/roundtrip.go b/internal/integration/roundtrip.go
index c9343d5..a020d7b 100644
--- a/internal/integration/roundtrip.go
+++ b/internal/integration/roundtrip.go
@@ -156,7 +156,7 @@ func (md *MyDialer) DialContext(ctx context.Context, urlStr string, requestHeade
 	return graphql.WSConn(conn), err
 }
 
-func newRoundtripWebScoketClient(t *testing.T, endpoint string) graphql.WebSocketClient {
+func newRoundtripWebSocketClient(t *testing.T, endpoint string) graphql.WebSocketClient {
 	dialer := websocket.DefaultDialer
 	if !strings.HasPrefix(endpoint, "ws") {
 		_, address, _ := strings.Cut(endpoint, "://")
diff --git a/internal/integration/server/server.go b/internal/integration/server/server.go
index 6526825..2c79909 100644
--- a/internal/integration/server/server.go
+++ b/internal/integration/server/server.go
@@ -159,11 +159,15 @@ func (m mutationResolver) CreateUser(ctx context.Context, input NewUser) (*User,
 func (s *subscriptionResolver) Count(ctx context.Context) (<-chan int, error) {
 	respChan := make(chan int, 1)
 	go func(respChan chan int) {
+		defer close(respChan)
 		counter := 0
 		for {
+			if counter == 10 {
+				return
+			}
 			respChan <- counter
 			counter++
-			time.Sleep(time.Second)
+			time.Sleep(100 * time.Millisecond)
 		}
 	}(respChan)
 	return respChan, nil


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 301ade2..96bc0b5 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -57,6 +57,11 @@ func TestMutation(t *testing.T) {
 	require.Errorf(t, err, "client does not support mutations")
 }
 
+type subscriptionResult struct {
+	clientUnsubscribed  bool
+	serverChannelClosed bool
+}
+
 func TestSubscription(t *testing.T) {
 	_ = `# @genqlient
 	subscription count { count }`
@@ -64,37 +69,78 @@ func TestSubscription(t *testing.T) {
 	ctx := context.Background()
 	server := server.RunServer()
 	defer server.Close()
-	wsClient := newRoundtripWebScoketClient(t, server.URL)
 
-	errChan, err := wsClient.Start(ctx)
-	require.NoError(t, err)
+	cases := []struct {
+		name           string
+		unsubThreshold time.Duration
+		expected       subscriptionResult
+	}{
+		{
+			name:           "server_closed_channel",
+			unsubThreshold: 5 * time.Second,
+			expected: subscriptionResult{
+				clientUnsubscribed:  false,
+				serverChannelClosed: true,
+			},
+		},
+		{
+			name:           "client_unsubscribed",
+			unsubThreshold: 300 * time.Millisecond,
+			expected: subscriptionResult{
+				clientUnsubscribed:  true,
+				serverChannelClosed: false,
+			},
+		},
+	}
 
-	dataChan, subscriptionID, err := count(ctx, wsClient)
-	require.NoError(t, err)
-	defer wsClient.Close()
-	counter := 0
-	start := time.Now()
-	for loop := true; loop; {
-		select {
-		case resp, more := <-dataChan:
-			if !more {
-				loop = false
-				break
-			}
-			require.NotNil(t, resp.Data)
-			assert.Equal(t, counter, resp.Data.Count)
-			require.Nil(t, resp.Errors)
-			if time.Since(start) > time.Second*5 {
-				err = wsClient.Unsubscribe(subscriptionID)
-				require.NoError(t, err)
-				loop = false
-			}
-			counter++
-		case err := <-errChan:
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			wsClient := newRoundtripWebSocketClient(t, server.URL)
+			errChan, err := wsClient.Start(ctx)
 			require.NoError(t, err)
-		case <-time.After(time.Second * 10):
-			require.NoError(t, fmt.Errorf("subscription timed out"))
-		}
+
+			dataChan, subscriptionID, err := count(ctx, wsClient)
+			require.NoError(t, err)
+			defer wsClient.Close()
+
+			var (
+				counter = 0
+				start   = time.Now()
+				result  = subscriptionResult{}
+			)
+
+			for loop := true; loop; {
+				select {
+				case resp, more := <-dataChan:
+					if !more {
+						result.serverChannelClosed = true
+						loop = false
+						break
+					}
+
+					require.NotNil(t, resp.Data)
+					assert.Equal(t, counter, resp.Data.Count)
+					require.Nil(t, resp.Errors)
+
+					if time.Since(start) > tc.unsubThreshold {
+						err := wsClient.Unsubscribe(subscriptionID)
+						require.NoError(t, err)
+						result.clientUnsubscribed = true
+						loop = false
+					}
+
+					counter++
+
+				case err := <-errChan:
+					require.NoError(t, err)
+
+				case <-time.After(10 * time.Second):
+					require.NoError(t, fmt.Errorf("subscription timed out"))
+				}
+			}
+
+			assert.Equal(t, tc.expected, result)
+		})
 	}
 }