Merge pull request #339 from 99designs/fix-subscription-goroutine-leak

Fix gouroutine leak when using subscriptions

client/websocket.go

@@ -14,6 +14,7 @@ const (
 	startMsg          = "start"           // Client -> Server
 	connectionAckMsg  = "connection_ack"  // Server -> Client
 	dataMsg           = "data"            // Server -> Client
+	errorMsg          = "error"           // Server -> Client
 )
 
 type operationMessage struct {
@@ -73,7 +74,11 @@ func (p *Client) Websocket(query string, options ...Option) *Subscription {
 			var op operationMessage
 			c.ReadJSON(&op)
 			if op.Type != dataMsg {
-				return fmt.Errorf("expected data message, got %#v", op)
+				if op.Type == errorMsg {
+					return fmt.Errorf(string(op.Payload))
+				} else {
+					return fmt.Errorf("expected data message, got %#v", op)
+				}
 			}
 
 			respDataRaw := map[string]interface{}{}

codegen/templates/generated.gotpl

@@ -176,6 +176,10 @@ func (e *executableSchema) Subscription(ctx context.Context, op *ast.OperationDe
 				return buf.Bytes()
 			})
 
+			if buf == nil {
+				return nil
+			}
+
 			return &graphql.Response{
 				Data:   buf,
 				Errors: ec.Errors,

codegen/testserver/generated_test.go

@@ -8,7 +8,9 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"runtime"
 	"testing"
+	"time"
 
 	"github.com/99designs/gqlgen/client"
 	"github.com/99designs/gqlgen/handler"
@@ -28,8 +30,9 @@ func TestForcedResolverFieldIsPointer(t *testing.T) {
 }
 
 func TestGeneratedServer(t *testing.T) {
+	tickChan := make(chan string, 1)
 	srv := httptest.NewServer(handler.GraphQL(NewExecutableSchema(Config{
-		Resolvers: &testResolver{},
+		Resolvers: &testResolver{tick: tickChan},
 	})))
 	c := client.New(srv.URL)
 
@@ -77,10 +80,39 @@ func TestGeneratedServer(t *testing.T) {
 			require.Nil(t, resp.ErrorBubble)
 			require.Equal(t, "Ok", resp.Valid)
 		})
+
+	})
+
+	t.Run("subscriptions", func(t *testing.T) {
+		t.Run("wont leak goroutines", func(t *testing.T) {
+			initialGoroutineCount := runtime.NumGoroutine()
+
+			sub := c.Websocket(`subscription { updated }`)
+
+			tickChan <- "message"
+
+			var msg struct {
+				resp struct {
+					Updated string
+				}
+			}
+
+			err := sub.Next(&msg.resp)
+			require.NoError(t, err)
+			require.Equal(t, "message", msg.resp.Updated)
+			sub.Close()
+
+			// need a little bit of time for goroutines to settle
+			time.Sleep(200 * time.Millisecond)
+
+			require.Equal(t, initialGoroutineCount, runtime.NumGoroutine())
+		})
 	})
 }
 
-type testResolver struct{}
+type testResolver struct {
+	tick chan string
+}
 
 func (r *testResolver) ForcedResolver() ForcedResolverResolver {
 	return &forcedResolverResolver{nil}
@@ -98,3 +130,26 @@ func (r *testQueryResolver) ErrorBubble(ctx context.Context) (*Error, error) {
 func (r *testQueryResolver) Valid(ctx context.Context) (string, error) {
 	return "Ok", nil
 }
+
+func (r *testResolver) Subscription() SubscriptionResolver {
+	return &testSubscriptionResolver{r}
+}
+
+type testSubscriptionResolver struct{ *testResolver }
+
+func (r *testSubscriptionResolver) Updated(ctx context.Context) (<-chan string, error) {
+	res := make(chan string, 1)
+
+	go func() {
+		for {
+			select {
+			case t := <-r.tick:
+				res <- t
+			case <-ctx.Done():
+				close(res)
+				return
+			}
+		}
+	}()
+	return res, nil
+}

codegen/testserver/resolver.go

@@ -17,6 +17,9 @@ func (r *Resolver) ForcedResolver() ForcedResolverResolver {
 func (r *Resolver) Query() QueryResolver {
 	return &queryResolver{r}
 }
+func (r *Resolver) Subscription() SubscriptionResolver {
+	return &subscriptionResolver{r}
+}
 
 type forcedResolverResolver struct{ *Resolver }
 
@@ -59,3 +62,9 @@ func (r *queryResolver) Valid(ctx context.Context) (string, error) {
 func (r *queryResolver) KeywordArgs(ctx context.Context, breakArg string, defaultArg string, funcArg string, interfaceArg string, selectArg string, caseArg string, deferArg string, goArg string, mapArg string, structArg string, chanArg string, elseArg string, gotoArg string, packageArg string, switchArg string, constArg string, fallthroughArg string, ifArg string, rangeArg string, typeArg string, continueArg string, forArg string, importArg string, returnArg string, varArg string) (bool, error) {
 	panic("not implemented")
 }
+
+type subscriptionResolver struct{ *Resolver }
+
+func (r *subscriptionResolver) Updated(ctx context.Context) (<-chan string, error) {
+	panic("not implemented")
+}

codegen/testserver/schema.graphql

@@ -11,6 +11,10 @@ type Query {
     valid: String!
 }
 
+type Subscription {
+    updated: String!
+}
+
 type Error {
     id: ID!
     errorOnNonRequiredField: String