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
5package trace
6
7import (
8	"cmd/internal/traceviewer"
9	"context"
10	"encoding/json"
11	"errors"
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.
21const (
22	phaseDurationBegin = "B"
23	phaseDurationEnd   = "E"
24	phaseFlowStart     = "s"
25	phaseFlowEnd       = "f"
26
27	bindEnclosingSlice = "e"
28)
29
30var traceStarted int32
31
32func getTraceContext(ctx context.Context) (traceContext, bool) {
33	if atomic.LoadInt32(&traceStarted) == 0 {
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.
44func 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(&traceviewer.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.
63func 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.
73func 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(&traceviewer.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(&traceviewer.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
99type Span struct {
100	t *tracer
101
102	name  string
103	tid   uint64
104	start time.Time
105	end   time.Time
106}
107
108func (s *Span) Done() {
109	if s == nil {
110		return
111	}
112	s.end = time.Now()
113	s.t.writeEvent(&traceviewer.Event{
114		Name:  s.name,
115		Time:  float64(s.end.UnixNano()) / float64(time.Microsecond),
116		TID:   s.tid,
117		Phase: phaseDurationEnd,
118	})
119}
120
121type tracer struct {
122	file chan traceFile // 1-buffered
123
124	nextTID    uint64
125	nextFlowID uint64
126}
127
128func (t *tracer) writeEvent(ev *traceviewer.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
152func (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
163func (t *tracer) getNextTID() uint64 {
164	return atomic.AddUint64(&t.nextTID, 1)
165}
166
167func (t *tracer) getNextFlowID() uint64 {
168	return atomic.AddUint64(&t.nextFlowID, 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.
173type traceKey struct{}
174
175type traceContext struct {
176	t   *tracer
177	tid uint64
178}
179
180// Start starts a trace which writes to the given file.
181func Start(ctx context.Context, file string) (context.Context, func() error, error) {
182	atomic.StoreInt32(&traceStarted, 1)
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
201type traceFile struct {
202	f       *os.File
203	sb      *strings.Builder
204	enc     *json.Encoder
205	entries int64
206}
207