github.com/v2fly/tools@v0.100.0/internal/event/export/trace.go (about)

     1  // Copyright 2019 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package export
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"sync"
    11  
    12  	"github.com/v2fly/tools/internal/event"
    13  	"github.com/v2fly/tools/internal/event/core"
    14  	"github.com/v2fly/tools/internal/event/keys"
    15  	"github.com/v2fly/tools/internal/event/label"
    16  )
    17  
    18  type SpanContext struct {
    19  	TraceID TraceID
    20  	SpanID  SpanID
    21  }
    22  
    23  type Span struct {
    24  	Name     string
    25  	ID       SpanContext
    26  	ParentID SpanID
    27  	mu       sync.Mutex
    28  	start    core.Event
    29  	finish   core.Event
    30  	events   []core.Event
    31  }
    32  
    33  type contextKeyType int
    34  
    35  const (
    36  	spanContextKey = contextKeyType(iota)
    37  	labelContextKey
    38  )
    39  
    40  func GetSpan(ctx context.Context) *Span {
    41  	v := ctx.Value(spanContextKey)
    42  	if v == nil {
    43  		return nil
    44  	}
    45  	return v.(*Span)
    46  }
    47  
    48  // Spans creates an exporter that maintains hierarchical span structure in the
    49  // context.
    50  // It creates new spans on start events, adds events to the current span on
    51  // log or label, and closes the span on end events.
    52  // The span structure can then be used by other exporters.
    53  func Spans(output event.Exporter) event.Exporter {
    54  	return func(ctx context.Context, ev core.Event, lm label.Map) context.Context {
    55  		switch {
    56  		case event.IsLog(ev), event.IsLabel(ev):
    57  			if span := GetSpan(ctx); span != nil {
    58  				span.mu.Lock()
    59  				span.events = append(span.events, ev)
    60  				span.mu.Unlock()
    61  			}
    62  		case event.IsStart(ev):
    63  			span := &Span{
    64  				Name:  keys.Start.Get(lm),
    65  				start: ev,
    66  			}
    67  			if parent := GetSpan(ctx); parent != nil {
    68  				span.ID.TraceID = parent.ID.TraceID
    69  				span.ParentID = parent.ID.SpanID
    70  			} else {
    71  				span.ID.TraceID = newTraceID()
    72  			}
    73  			span.ID.SpanID = newSpanID()
    74  			ctx = context.WithValue(ctx, spanContextKey, span)
    75  		case event.IsEnd(ev):
    76  			if span := GetSpan(ctx); span != nil {
    77  				span.mu.Lock()
    78  				span.finish = ev
    79  				span.mu.Unlock()
    80  			}
    81  		case event.IsDetach(ev):
    82  			ctx = context.WithValue(ctx, spanContextKey, nil)
    83  		}
    84  		return output(ctx, ev, lm)
    85  	}
    86  }
    87  
    88  func (s *SpanContext) Format(f fmt.State, r rune) {
    89  	fmt.Fprintf(f, "%v:%v", s.TraceID, s.SpanID)
    90  }
    91  
    92  func (s *Span) Start() core.Event {
    93  	// start never changes after construction, so we dont need to hold the mutex
    94  	return s.start
    95  }
    96  
    97  func (s *Span) Finish() core.Event {
    98  	s.mu.Lock()
    99  	defer s.mu.Unlock()
   100  	return s.finish
   101  }
   102  
   103  func (s *Span) Events() []core.Event {
   104  	s.mu.Lock()
   105  	defer s.mu.Unlock()
   106  	return s.events
   107  }
   108  
   109  func (s *Span) Format(f fmt.State, r rune) {
   110  	s.mu.Lock()
   111  	defer s.mu.Unlock()
   112  	fmt.Fprintf(f, "%v %v", s.Name, s.ID)
   113  	if s.ParentID.IsValid() {
   114  		fmt.Fprintf(f, "[%v]", s.ParentID)
   115  	}
   116  	fmt.Fprintf(f, " %v->%v", s.start, s.finish)
   117  }