github.com/senomas/gqlgen@v0.17.11-0.20220626120754-9aee61b0716a/graphql/handler/apollofederatedtracingv1/tree_builder.go (about)

     1  package apollofederatedtracingv1
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/99designs/gqlgen/graphql"
    10  	"github.com/99designs/gqlgen/graphql/handler/apollofederatedtracingv1/generated"
    11  	"google.golang.org/protobuf/types/known/timestamppb"
    12  )
    13  
    14  type TreeBuilder struct {
    15  	Trace    *generated.Trace
    16  	rootNode generated.Trace_Node
    17  	nodes    map[string]NodeMap // nodes is used to store a pointer map using the node path (e.g. todo[0].id) to itself as well as it's parent
    18  
    19  	startTime *time.Time
    20  	stopped   bool
    21  	mu        sync.Mutex
    22  }
    23  
    24  type NodeMap struct {
    25  	self   *generated.Trace_Node
    26  	parent *generated.Trace_Node
    27  }
    28  
    29  // NewTreeBuilder is used to start the node tree with a default root node, along with the related tree nodes map entry
    30  func NewTreeBuilder() *TreeBuilder {
    31  	tb := TreeBuilder{
    32  		rootNode: generated.Trace_Node{},
    33  	}
    34  
    35  	t := generated.Trace{
    36  		Root: &tb.rootNode,
    37  	}
    38  	tb.nodes = make(map[string]NodeMap)
    39  	tb.nodes[""] = NodeMap{self: &tb.rootNode, parent: nil}
    40  
    41  	tb.Trace = &t
    42  
    43  	return &tb
    44  }
    45  
    46  // StartTimer marks the time using protobuf timestamp format for use in timing calculations
    47  func (tb *TreeBuilder) StartTimer(ctx context.Context) {
    48  	if tb.startTime != nil {
    49  		fmt.Println(fmt.Errorf("StartTimer called twice"))
    50  	}
    51  	if tb.stopped {
    52  		fmt.Println(fmt.Errorf("StartTimer called after StopTimer"))
    53  	}
    54  
    55  	rc := graphql.GetOperationContext(ctx)
    56  	start := rc.Stats.OperationStart
    57  
    58  	tb.Trace.StartTime = timestamppb.New(start)
    59  	tb.startTime = &start
    60  }
    61  
    62  // StopTimer marks the end of the timer, along with setting the related fields in the protobuf representation
    63  func (tb *TreeBuilder) StopTimer(ctx context.Context) {
    64  	if tb.startTime == nil {
    65  		fmt.Println(fmt.Errorf("StopTimer called before StartTimer"))
    66  	}
    67  	if tb.stopped {
    68  		fmt.Println(fmt.Errorf("StopTimer called twice"))
    69  	}
    70  
    71  	ts := graphql.Now().UTC()
    72  	tb.Trace.DurationNs = uint64(ts.Sub(*tb.startTime).Nanoseconds())
    73  	tb.Trace.EndTime = timestamppb.New(ts)
    74  	tb.stopped = true
    75  }
    76  
    77  // On each field, it calculates the time started at as now - tree.StartTime, as well as a deferred function upon full resolution of the
    78  // field as now - tree.StartTime; these are used by Apollo to calculate how fields are being resolved in the AST
    79  func (tb *TreeBuilder) WillResolveField(ctx context.Context) {
    80  	if tb.startTime == nil {
    81  		fmt.Println(fmt.Errorf("WillResolveField called before StartTimer"))
    82  		return
    83  	}
    84  	if tb.stopped {
    85  		fmt.Println(fmt.Errorf("WillResolveField called after StopTimer"))
    86  		return
    87  	}
    88  	fc := graphql.GetFieldContext(ctx)
    89  
    90  	node := tb.newNode(fc)
    91  	node.StartTime = uint64(graphql.Now().Sub(*tb.startTime).Nanoseconds())
    92  	defer func() {
    93  		node.EndTime = uint64(graphql.Now().Sub(*tb.startTime).Nanoseconds())
    94  	}()
    95  
    96  	node.Type = fc.Field.Definition.Type.String()
    97  	node.ParentType = fc.Object
    98  }
    99  
   100  // newNode is called on each new node within the AST and sets related values such as the entry in the tree.node map and ID attribute
   101  func (tb *TreeBuilder) newNode(path *graphql.FieldContext) *generated.Trace_Node {
   102  	// if the path is empty, it is the root node of the operation
   103  	if path.Path().String() == "" {
   104  		return &tb.rootNode
   105  	}
   106  
   107  	self := &generated.Trace_Node{}
   108  	pn := tb.ensureParentNode(path)
   109  
   110  	if path.Index != nil {
   111  		self.Id = &generated.Trace_Node_Index{Index: uint32(*path.Index)}
   112  	} else {
   113  		self.Id = &generated.Trace_Node_ResponseName{ResponseName: path.Field.Name}
   114  	}
   115  
   116  	// lock the map from being read/written concurrently to avoid panics
   117  	tb.mu.Lock()
   118  	nodeRef := tb.nodes[path.Path().String()]
   119  	// set the values for the node references to help build the tree
   120  	nodeRef.parent = pn
   121  	nodeRef.self = self
   122  
   123  	// since they are references, we point the parent to it's children nodes
   124  	nodeRef.parent.Child = append(nodeRef.parent.Child, self)
   125  	nodeRef.self = self
   126  	tb.nodes[path.Path().String()] = nodeRef
   127  	tb.mu.Unlock()
   128  
   129  	return self
   130  }
   131  
   132  // ensureParentNode ensures the node isn't orphaned
   133  func (tb *TreeBuilder) ensureParentNode(path *graphql.FieldContext) *generated.Trace_Node {
   134  	// lock to read briefly, then unlock to avoid r/w issues
   135  	tb.mu.Lock()
   136  	nodeRef := tb.nodes[path.Parent.Path().String()]
   137  	tb.mu.Unlock()
   138  
   139  	if nodeRef.self != nil {
   140  		return nodeRef.self
   141  	}
   142  
   143  	return tb.newNode(path.Parent)
   144  }