ensure subscribe and resolve use the same cancellable context

Author: branden-blackline

executor.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"sort"
 	"strings"
 
 	"github.com/graphql-go/graphql/gqlerrors"
@@ -254,7 +255,9 @@ func executeFieldsSerially(p executeFieldsParams) *Result {
 	}
 
 	finalResults := make(map[string]interface{}, len(p.Fields))
-	for responseName, fieldASTs := range p.Fields {
+	for _, orderedField := range orderedFields(p.Fields) {
+		responseName := orderedField.responseName
+		fieldASTs := orderedField.fieldASTs
 		fieldPath := p.Path.WithKey(responseName)
 		resolved, state := resolveField(p.ExecutionContext, p.ParentType, p.Source, fieldASTs, fieldPath)
 		if state.hasNoFieldDefs {
@@ -1038,3 +1041,39 @@ func getFieldDef(schema Schema, parentType *Object, fieldName string) *FieldDefi
 	}
 	return parentType.Fields()[fieldName]
 }
+
+// contains field information that will be placed in an ordered slice
+type orderedField struct {
+	responseName string
+	fieldASTs    []*ast.Field
+}
+
+// orders fields from a fields map by location in the source
+func orderedFields(fields map[string][]*ast.Field) []*orderedField {
+	orderedFields := []*orderedField{}
+	fieldMap := map[int]*orderedField{}
+	startLocs := []int{}
+
+	for responseName, fieldASTs := range fields {
+		// find the lowest location in the current fieldASTs
+		lowest := -1
+		for _, fieldAST := range fieldASTs {
+			loc := fieldAST.GetLoc().Start
+			if lowest == -1 || loc < lowest {
+				lowest = loc
+			}
+		}
+		startLocs = append(startLocs, lowest)
+		fieldMap[lowest] = &orderedField{
+			responseName: responseName,
+			fieldASTs:    fieldASTs,
+		}
+	}
+
+	sort.Ints(startLocs)
+	for _, startLoc := range startLocs {
+		orderedFields = append(orderedFields, fieldMap[startLoc])
+	}
+
+	return orderedFields
+}

subscription.go

@@ -39,18 +39,13 @@ type ResultIterator struct {
 }
 
 // NewResultIterator creates a new iterator and starts handling message on the result channel
-func NewResultIterator(ctx context.Context, ch chan *Result) *ResultIterator {
-	if ctx == nil {
-		ctx = context.Background()
-	}
-
-	cctx, cancelFunc := context.WithCancel(ctx)
+func NewResultIterator(ctx context.Context, cancelFunc context.CancelFunc, ch chan *Result) *ResultIterator {
 	iterator := &ResultIterator{
 		currentHandlerID: 0,
 		count:            0,
-		ctx:              cctx,
-		ch:               ch,
+		ctx:              ctx,
 		cancelFunc:       cancelFunc,
+		ch:               ch,
 		cancelled:        false,
 		handlers:         map[int64]*subscriptionHanlderConfig{},
 	}
@@ -140,14 +135,16 @@ func Subscribe(p SubscribeParams) *ResultIterator {
 		ctx = context.Background()
 	}
 
+	sctx, cancelFunc := context.WithCancel(ctx)
+
 	var mapSourceToResponse = func(payload interface{}) *Result {
 		return Execute(ExecuteParams{
 			Schema:        p.Schema,
 			Root:          payload,
 			AST:           p.Document,
 			OperationName: p.OperationName,
 			Args:          p.VariableValues,
-			Context:       p.ContextValue,
+			Context:       sctx,
 		})
 	}
 
@@ -168,7 +165,7 @@ func Subscribe(p SubscribeParams) *ResultIterator {
 			OperationName: p.OperationName,
 			Args:          p.VariableValues,
 			Result:        result,
-			Context:       p.ContextValue,
+			Context:       sctx,
 		})
 
 		if err != nil {
@@ -236,7 +233,7 @@ func Subscribe(p SubscribeParams) *ResultIterator {
 			Source:  p.RootValue,
 			Args:    args,
 			Info:    info,
-			Context: exeContext.Context,
+			Context: sctx,
 		})
 		if err != nil {
 			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
@@ -255,11 +252,9 @@ func Subscribe(p SubscribeParams) *ResultIterator {
 		case chan interface{}:
 			for {
 				select {
-				case <-ctx.Done():
-					fmt.Printf("done context called")
+				case <-sctx.Done():
 					return
 				case res := <-fieldResult.(chan interface{}):
-
 					resultChannel <- mapSourceToResponse(res)
 				}
 			}
@@ -270,5 +265,5 @@ func Subscribe(p SubscribeParams) *ResultIterator {
 	}()
 
 	// return a result iterator
-	return NewResultIterator(p.ContextValue, resultChannel)
+	return NewResultIterator(sctx, cancelFunc, resultChannel)
 }