github.com/niko0xdev/gqlgen@v0.17.55-0.20240120102243-2ecff98c3e37/graphql/handler/apollofederatedtracingv1/tree_builder.go (about)

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