trace.go

     1  // Copyright 2020 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 trace
     6  
     7  import (
     8  	"context"
     9  	"encoding/json"
    10  	"errors"
    11  	"internal/trace/traceviewer/format"
    12  	"os"
    13  	"strings"
    14  	"sync/atomic"
    15  	"time"
    16  )
    17  
    18  // Constants used in event fields.
    19  // See https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU
    20  // for more details.
    21  const (
    22  	phaseDurationBegin = "B"
    23  	phaseDurationEnd   = "E"
    24  	phaseFlowStart     = "s"
    25  	phaseFlowEnd       = "f"
    26  
    27  	bindEnclosingSlice = "e"
    28  )
    29  
    30  var traceStarted atomic.Bool
    31  
    32  func getTraceContext(ctx context.Context) (traceContext, bool) {
    33  	if !traceStarted.Load() {
    34  		return traceContext{}, false
    35  	}
    36  	v := ctx.Value(traceKey{})
    37  	if v == nil {
    38  		return traceContext{}, false
    39  	}
    40  	return v.(traceContext), true
    41  }
    42  
    43  // StartSpan starts a trace event with the given name. The Span ends when its Done method is called.
    44  func StartSpan(ctx context.Context, name string) (context.Context, *Span) {
    45  	tc, ok := getTraceContext(ctx)
    46  	if !ok {
    47  		return ctx, nil
    48  	}
    49  	childSpan := &Span{t: tc.t, name: name, tid: tc.tid, start: time.Now()}
    50  	tc.t.writeEvent(&format.Event{
    51  		Name:  childSpan.name,
    52  		Time:  float64(childSpan.start.UnixNano()) / float64(time.Microsecond),
    53  		TID:   childSpan.tid,
    54  		Phase: phaseDurationBegin,
    55  	})
    56  	ctx = context.WithValue(ctx, traceKey{}, traceContext{tc.t, tc.tid})
    57  	return ctx, childSpan
    58  }
    59  
    60  // StartGoroutine associates the context with a new Thread ID. The Chrome trace viewer associates each
    61  // trace event with a thread, and doesn't expect events with the same thread id to happen at the
    62  // same time.
    63  func StartGoroutine(ctx context.Context) context.Context {
    64  	tc, ok := getTraceContext(ctx)
    65  	if !ok {
    66  		return ctx
    67  	}
    68  	return context.WithValue(ctx, traceKey{}, traceContext{tc.t, tc.t.getNextTID()})
    69  }
    70  
    71  // Flow marks a flow indicating that the 'to' span depends on the 'from' span.
    72  // Flow should be called while the 'to' span is in progress.
    73  func Flow(ctx context.Context, from *Span, to *Span) {
    74  	tc, ok := getTraceContext(ctx)
    75  	if !ok || from == nil || to == nil {
    76  		return
    77  	}
    78  
    79  	id := tc.t.getNextFlowID()
    80  	tc.t.writeEvent(&format.Event{
    81  		Name:     from.name + " -> " + to.name,
    82  		Category: "flow",
    83  		ID:       id,
    84  		Time:     float64(from.end.UnixNano()) / float64(time.Microsecond),
    85  		Phase:    phaseFlowStart,
    86  		TID:      from.tid,
    87  	})
    88  	tc.t.writeEvent(&format.Event{
    89  		Name:      from.name + " -> " + to.name,
    90  		Category:  "flow", // TODO(matloob): Add Category to Flow?
    91  		ID:        id,
    92  		Time:      float64(to.start.UnixNano()) / float64(time.Microsecond),
    93  		Phase:     phaseFlowEnd,
    94  		TID:       to.tid,
    95  		BindPoint: bindEnclosingSlice,
    96  	})
    97  }
    98  
    99  type Span struct {
   100  	t *tracer
   101  
   102  	name  string
   103  	tid   uint64
   104  	start time.Time
   105  	end   time.Time
   106  }
   107  
   108  func (s *Span) Done() {
   109  	if s == nil {
   110  		return
   111  	}
   112  	s.end = time.Now()
   113  	s.t.writeEvent(&format.Event{
   114  		Name:  s.name,
   115  		Time:  float64(s.end.UnixNano()) / float64(time.Microsecond),
   116  		TID:   s.tid,
   117  		Phase: phaseDurationEnd,
   118  	})
   119  }
   120  
   121  type tracer struct {
   122  	file chan traceFile // 1-buffered
   123  
   124  	nextTID    atomic.Uint64
   125  	nextFlowID atomic.Uint64
   126  }
   127  
   128  func (t *tracer) writeEvent(ev *format.Event) error {
   129  	f := <-t.file
   130  	defer func() { t.file <- f }()
   131  	var err error
   132  	if f.entries == 0 {
   133  		_, err = f.sb.WriteString("[\n")
   134  	} else {
   135  		_, err = f.sb.WriteString(",")
   136  	}
   137  	f.entries++
   138  	if err != nil {
   139  		return nil
   140  	}
   141  
   142  	if err := f.enc.Encode(ev); err != nil {
   143  		return err
   144  	}
   145  
   146  	// Write event string to output file.
   147  	_, err = f.f.WriteString(f.sb.String())
   148  	f.sb.Reset()
   149  	return err
   150  }
   151  
   152  func (t *tracer) Close() error {
   153  	f := <-t.file
   154  	defer func() { t.file <- f }()
   155  
   156  	_, firstErr := f.f.WriteString("]")
   157  	if err := f.f.Close(); firstErr == nil {
   158  		firstErr = err
   159  	}
   160  	return firstErr
   161  }
   162  
   163  func (t *tracer) getNextTID() uint64 {
   164  	return t.nextTID.Add(1)
   165  }
   166  
   167  func (t *tracer) getNextFlowID() uint64 {
   168  	return t.nextFlowID.Add(1)
   169  }
   170  
   171  // traceKey is the context key for tracing information. It is unexported to prevent collisions with context keys defined in
   172  // other packages.
   173  type traceKey struct{}
   174  
   175  type traceContext struct {
   176  	t   *tracer
   177  	tid uint64
   178  }
   179  
   180  // Start starts a trace which writes to the given file.
   181  func Start(ctx context.Context, file string) (context.Context, func() error, error) {
   182  	traceStarted.Store(true)
   183  	if file == "" {
   184  		return nil, nil, errors.New("no trace file supplied")
   185  	}
   186  	f, err := os.Create(file)
   187  	if err != nil {
   188  		return nil, nil, err
   189  	}
   190  	t := &tracer{file: make(chan traceFile, 1)}
   191  	sb := new(strings.Builder)
   192  	t.file <- traceFile{
   193  		f:   f,
   194  		sb:  sb,
   195  		enc: json.NewEncoder(sb),
   196  	}
   197  	ctx = context.WithValue(ctx, traceKey{}, traceContext{t: t})
   198  	return ctx, t.Close, nil
   199  }
   200  
   201  type traceFile struct {
   202  	f       *os.File
   203  	sb      *strings.Builder
   204  	enc     *json.Encoder
   205  	entries int64
   206  }
   207
View as plain text