Handle websocket data channel closing

Khan · Sep 27, 2024 · 9473b93 · 9473b93
1 parent 4466fc1
commit 9473b93
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 1 deletion.
diff --git a/generate/operation.go.tmpl b/generate/operation.go.tmpl
@@ -62,6 +62,14 @@ type {{.Name}}WsResponse struct {
 func {{.Name}}ForwardData(interfaceChan interface{}, jsonRawMsg json.RawMessage) error {
     var gqlResp graphql.Response
 	var wsResp {{.Name}}WsResponse
+	if len(jsonRawMsg) == 0 {
+		dataChan_, ok := interfaceChan.(chan {{.Name}}WsResponse)
+		if !ok {
+			return errors.New("failed to cast interface into 'chan {{.Name}}WsResponse'")
+		}
+		close(dataChan_)
+		return nil
+	}
     err := json.Unmarshal(jsonRawMsg, &gqlResp)
     if err != nil {
         return err

diff --git a/internal/integration/generated.go b/internal/integration/generated.go
diff --git 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 == 3 {
+				return
+			}
 			respChan <- counter
 			counter++
-			time.Sleep(time.Second)
+			time.Sleep(10 * time.Millisecond)
 		}
 	}(respChan)
 	return respChan, nil