139 lines
3.0 KiB
Go
139 lines
3.0 KiB
Go
package tracing
|
|
|
|
import (
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// The Trace type functions as a container for capturing Spans used to
|
|
// trace the execution of a request.
|
|
type Trace struct {
|
|
mu sync.Mutex
|
|
spans map[uint64]RawSpan
|
|
}
|
|
|
|
// NewTrace starts a new trace and returns a root span identified by the provided name.
|
|
//
|
|
// Additional options may be specified to override the default behavior when creating the span.
|
|
func NewTrace(name string, opt ...StartSpanOption) (*Trace, *Span) {
|
|
t := &Trace{spans: make(map[uint64]RawSpan)}
|
|
s := &Span{tracer: t}
|
|
s.raw.Name = name
|
|
s.raw.Context.TraceID, s.raw.Context.SpanID = randomID2()
|
|
setOptions(s, opt)
|
|
|
|
return t, s
|
|
}
|
|
|
|
// NewTraceFromSpan starts a new trace and returns the associated span, which is a child of the
|
|
// parent span context.
|
|
func NewTraceFromSpan(name string, parent SpanContext, opt ...StartSpanOption) (*Trace, *Span) {
|
|
t := &Trace{spans: make(map[uint64]RawSpan)}
|
|
s := &Span{tracer: t}
|
|
s.raw.Name = name
|
|
s.raw.ParentSpanID = parent.SpanID
|
|
s.raw.Context.TraceID = parent.TraceID
|
|
s.raw.Context.SpanID = randomID()
|
|
setOptions(s, opt)
|
|
|
|
return t, s
|
|
}
|
|
|
|
func (t *Trace) startSpan(name string, sc SpanContext, opt []StartSpanOption) *Span {
|
|
s := &Span{tracer: t}
|
|
s.raw.Name = name
|
|
s.raw.Context.SpanID = randomID()
|
|
s.raw.Context.TraceID = sc.TraceID
|
|
s.raw.ParentSpanID = sc.SpanID
|
|
setOptions(s, opt)
|
|
|
|
return s
|
|
}
|
|
|
|
func setOptions(s *Span, opt []StartSpanOption) {
|
|
for _, o := range opt {
|
|
o.applyStart(s)
|
|
}
|
|
|
|
if s.raw.Start.IsZero() {
|
|
s.raw.Start = time.Now()
|
|
}
|
|
}
|
|
|
|
func (t *Trace) addRawSpan(raw RawSpan) {
|
|
t.mu.Lock()
|
|
t.spans[raw.Context.SpanID] = raw
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
// Tree returns a graph of the current trace.
|
|
func (t *Trace) Tree() *TreeNode {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
|
|
for _, s := range t.spans {
|
|
if s.ParentSpanID == 0 {
|
|
return t.treeFrom(s.Context.SpanID)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Merge combines other with the current trace. This is
|
|
// typically necessary when traces are transferred from a remote.
|
|
func (t *Trace) Merge(other *Trace) {
|
|
for k, s := range other.spans {
|
|
t.spans[k] = s
|
|
}
|
|
}
|
|
|
|
func (t *Trace) TreeFrom(root uint64) *TreeNode {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
return t.treeFrom(root)
|
|
}
|
|
|
|
func (t *Trace) treeFrom(root uint64) *TreeNode {
|
|
c := map[uint64]*TreeNode{}
|
|
|
|
for k, s := range t.spans {
|
|
c[k] = &TreeNode{Raw: s}
|
|
}
|
|
|
|
if _, ok := c[root]; !ok {
|
|
return nil
|
|
}
|
|
|
|
for _, n := range c {
|
|
if n.Raw.ParentSpanID != 0 {
|
|
if pn := c[n.Raw.ParentSpanID]; pn != nil {
|
|
pn.Children = append(pn.Children, n)
|
|
}
|
|
}
|
|
}
|
|
|
|
// sort nodes
|
|
var v treeSortVisitor
|
|
Walk(&v, c[root])
|
|
|
|
return c[root]
|
|
}
|
|
|
|
type treeSortVisitor struct{}
|
|
|
|
func (v *treeSortVisitor) Visit(node *TreeNode) Visitor {
|
|
sort.Slice(node.Children, func(i, j int) bool {
|
|
lt, rt := node.Children[i].Raw.Start.UnixNano(), node.Children[j].Raw.Start.UnixNano()
|
|
if lt < rt {
|
|
return true
|
|
} else if lt > rt {
|
|
return false
|
|
}
|
|
|
|
ln, rn := node.Children[i].Raw.Name, node.Children[j].Raw.Name
|
|
return ln < rn
|
|
})
|
|
return v
|
|
}
|