1// Copyright 2018 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	"context"
9	"fmt"
10	"sync/atomic"
11	_ "unsafe"
12)
13
14type traceContextKey struct{}
15
16// NewTask creates a task instance with the type taskType and returns
17// it along with a Context that carries the task.
18// If the input context contains a task, the new task is its subtask.
19//
20// The taskType is used to classify task instances. Analysis tools
21// like the Go execution tracer may assume there are only a bounded
22// number of unique task types in the system.
23//
24// The returned Task's [Task.End] method is used to mark the task's end.
25// The trace tool measures task latency as the time between task creation
26// and when the End method is called, and provides the latency
27// distribution per task type.
28// If the End method is called multiple times, only the first
29// call is used in the latency measurement.
30//
31//	ctx, task := trace.NewTask(ctx, "awesomeTask")
32//	trace.WithRegion(ctx, "preparation", prepWork)
33//	// preparation of the task
34//	go func() {  // continue processing the task in a separate goroutine.
35//	    defer task.End()
36//	    trace.WithRegion(ctx, "remainingWork", remainingWork)
37//	}()
38func NewTask(pctx context.Context, taskType string) (ctx context.Context, task *Task) {
39	pid := fromContext(pctx).id
40	id := newID()
41	userTaskCreate(id, pid, taskType)
42	s := &Task{id: id}
43	return context.WithValue(pctx, traceContextKey{}, s), s
44
45	// We allocate a new task even when
46	// the tracing is disabled because the context and task
47	// can be used across trace enable/disable boundaries,
48	// which complicates the problem.
49	//
50	// For example, consider the following scenario:
51	//   - trace is enabled.
52	//   - trace.WithRegion is called, so a new context ctx
53	//     with a new region is created.
54	//   - trace is disabled.
55	//   - trace is enabled again.
56	//   - trace APIs with the ctx is called. Is the ID in the task
57	//   a valid one to use?
58	//
59	// TODO(hyangah): reduce the overhead at least when
60	// tracing is disabled. Maybe the id can embed a tracing
61	// round number and ignore ids generated from previous
62	// tracing round.
63}
64
65func fromContext(ctx context.Context) *Task {
66	if s, ok := ctx.Value(traceContextKey{}).(*Task); ok {
67		return s
68	}
69	return &bgTask
70}
71
72// Task is a data type for tracing a user-defined, logical operation.
73type Task struct {
74	id uint64
75	// TODO(hyangah): record parent id?
76}
77
78// End marks the end of the operation represented by the [Task].
79func (t *Task) End() {
80	userTaskEnd(t.id)
81}
82
83var lastTaskID uint64 = 0 // task id issued last time
84
85func newID() uint64 {
86	// TODO(hyangah): use per-P cache
87	return atomic.AddUint64(&lastTaskID, 1)
88}
89
90var bgTask = Task{id: uint64(0)}
91
92// Log emits a one-off event with the given category and message.
93// Category can be empty and the API assumes there are only a handful of
94// unique categories in the system.
95func Log(ctx context.Context, category, message string) {
96	id := fromContext(ctx).id
97	userLog(id, category, message)
98}
99
100// Logf is like [Log], but the value is formatted using the specified format spec.
101func Logf(ctx context.Context, category, format string, args ...any) {
102	if IsEnabled() {
103		// Ideally this should be just Log, but that will
104		// add one more frame in the stack trace.
105		id := fromContext(ctx).id
106		userLog(id, category, fmt.Sprintf(format, args...))
107	}
108}
109
110const (
111	regionStartCode = uint64(0)
112	regionEndCode   = uint64(1)
113)
114
115// WithRegion starts a region associated with its calling goroutine, runs fn,
116// and then ends the region. If the context carries a task, the region is
117// associated with the task. Otherwise, the region is attached to the background
118// task.
119//
120// The regionType is used to classify regions, so there should be only a
121// handful of unique region types.
122func WithRegion(ctx context.Context, regionType string, fn func()) {
123	// NOTE:
124	// WithRegion helps avoiding misuse of the API but in practice,
125	// this is very restrictive:
126	// - Use of WithRegion makes the stack traces captured from
127	//   region start and end are identical.
128	// - Refactoring the existing code to use WithRegion is sometimes
129	//   hard and makes the code less readable.
130	//     e.g. code block nested deep in the loop with various
131	//          exit point with return values
132	// - Refactoring the code to use this API with closure can
133	//   cause different GC behavior such as retaining some parameters
134	//   longer.
135	// This causes more churns in code than I hoped, and sometimes
136	// makes the code less readable.
137
138	id := fromContext(ctx).id
139	userRegion(id, regionStartCode, regionType)
140	defer userRegion(id, regionEndCode, regionType)
141	fn()
142}
143
144// StartRegion starts a region and returns it.
145// The returned Region's [Region.End] method must be called
146// from the same goroutine where the region was started.
147// Within each goroutine, regions must nest. That is, regions started
148// after this region must be ended before this region can be ended.
149// Recommended usage is
150//
151//	defer trace.StartRegion(ctx, "myTracedRegion").End()
152func StartRegion(ctx context.Context, regionType string) *Region {
153	if !IsEnabled() {
154		return noopRegion
155	}
156	id := fromContext(ctx).id
157	userRegion(id, regionStartCode, regionType)
158	return &Region{id, regionType}
159}
160
161// Region is a region of code whose execution time interval is traced.
162type Region struct {
163	id         uint64
164	regionType string
165}
166
167var noopRegion = &Region{}
168
169// End marks the end of the traced code region.
170func (r *Region) End() {
171	if r == noopRegion {
172		return
173	}
174	userRegion(r.id, regionEndCode, r.regionType)
175}
176
177// IsEnabled reports whether tracing is enabled.
178// The information is advisory only. The tracing status
179// may have changed by the time this function returns.
180func IsEnabled() bool {
181	return tracing.enabled.Load()
182}
183
184//
185// Function bodies are defined in runtime/trace.go
186//
187
188// emits UserTaskCreate event.
189func userTaskCreate(id, parentID uint64, taskType string)
190
191// emits UserTaskEnd event.
192func userTaskEnd(id uint64)
193
194// emits UserRegion event.
195func userRegion(id, mode uint64, regionType string)
196
197// emits UserLog event.
198func userLog(id uint64, category, message string)
199