github.com/bir3/gocompiler@v0.9.2202/src/internal/trace/traceviewer/emitter.go

// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package traceviewer

import (
	"encoding/json"
	"fmt"
	"github.com/bir3/gocompiler/src/internal/trace"
	"github.com/bir3/gocompiler/src/internal/trace/traceviewer/format"
	"io"
	"strconv"
	"time"
)

type TraceConsumer struct {
	ConsumeTimeUnit		func(unit string)
	ConsumeViewerEvent	func(v *format.Event, required bool)
	ConsumeViewerFrame	func(key string, f format.Frame)
	Flush			func()
}

// ViewerDataTraceConsumer returns a TraceConsumer that writes to w. The
// startIdx and endIdx are used for splitting large traces. They refer to
// indexes in the traceEvents output array, not the events in the trace input.
func ViewerDataTraceConsumer(w io.Writer, startIdx, endIdx int64) TraceConsumer {
	allFrames := make(map[string]format.Frame)
	requiredFrames := make(map[string]format.Frame)
	enc := json.NewEncoder(w)
	written := 0
	index := int64(-1)

	io.WriteString(w, "{")
	return TraceConsumer{
		ConsumeTimeUnit: func(unit string) {
			io.WriteString(w, `"displayTimeUnit":`)
			enc.Encode(unit)
			io.WriteString(w, ",")
		},
		ConsumeViewerEvent: func(v *format.Event, required bool) {
			index++
			if !required && (index < startIdx || index > endIdx) {
				// not in the range. Skip!
				return
			}
			WalkStackFrames(allFrames, v.Stack, func(id int) {
				s := strconv.Itoa(id)
				requiredFrames[s] = allFrames[s]
			})
			WalkStackFrames(allFrames, v.EndStack, func(id int) {
				s := strconv.Itoa(id)
				requiredFrames[s] = allFrames[s]
			})
			if written == 0 {
				io.WriteString(w, `"traceEvents": [`)
			}
			if written > 0 {
				io.WriteString(w, ",")
			}
			enc.Encode(v)
			// TODO(mknyszek): get rid of the extra \n inserted by enc.Encode.
			// Same should be applied to splittingTraceConsumer.
			written++
		},
		ConsumeViewerFrame: func(k string, v format.Frame) {
			allFrames[k] = v
		},
		Flush: func() {
			io.WriteString(w, `], "stackFrames":`)
			enc.Encode(requiredFrames)
			io.WriteString(w, `}`)
		},
	}
}

func SplittingTraceConsumer(max int) (*splitter, TraceConsumer) {
	type eventSz struct {
		Time	float64
		Sz	int
		Frames	[]int
	}

	var (
		// data.Frames contains only the frames for required events.
		data	= format.Data{Frames: make(map[string]format.Frame)}

		allFrames	= make(map[string]format.Frame)

		sizes	[]eventSz
		cw	countingWriter
	)

	s := new(splitter)

	return s, TraceConsumer{
		ConsumeTimeUnit: func(unit string) {
			data.TimeUnit = unit
		},
		ConsumeViewerEvent: func(v *format.Event, required bool) {
			if required {
				// Store required events inside data so flush
				// can include them in the required part of the
				// trace.
				data.Events = append(data.Events, v)
				WalkStackFrames(allFrames, v.Stack, func(id int) {
					s := strconv.Itoa(id)
					data.Frames[s] = allFrames[s]
				})
				WalkStackFrames(allFrames, v.EndStack, func(id int) {
					s := strconv.Itoa(id)
					data.Frames[s] = allFrames[s]
				})
				return
			}
			enc := json.NewEncoder(&cw)
			enc.Encode(v)
			size := eventSz{Time: v.Time, Sz: cw.size + 1}	// +1 for ",".
			// Add referenced stack frames. Their size is computed
			// in flush, where we can dedup across events.
			WalkStackFrames(allFrames, v.Stack, func(id int) {
				size.Frames = append(size.Frames, id)
			})
			WalkStackFrames(allFrames, v.EndStack, func(id int) {
				size.Frames = append(size.Frames, id)	// This may add duplicates. We'll dedup later.
			})
			sizes = append(sizes, size)
			cw.size = 0
		},
		ConsumeViewerFrame: func(k string, v format.Frame) {
			allFrames[k] = v
		},
		Flush: func() {
			// Calculate size of the mandatory part of the trace.
			// This includes thread names and stack frames for
			// required events.
			cw.size = 0
			enc := json.NewEncoder(&cw)
			enc.Encode(data)
			requiredSize := cw.size

			// Then calculate size of each individual event and
			// their stack frames, grouping them into ranges. We
			// only include stack frames relevant to the events in
			// the range to reduce overhead.

			var (
				start	= 0

				eventsSize	= 0

				frames		= make(map[string]format.Frame)
				framesSize	= 0
			)
			for i, ev := range sizes {
				eventsSize += ev.Sz

				// Add required stack frames. Note that they
				// may already be in the map.
				for _, id := range ev.Frames {
					s := strconv.Itoa(id)
					_, ok := frames[s]
					if ok {
						continue
					}
					f := allFrames[s]
					frames[s] = f
					framesSize += stackFrameEncodedSize(uint(id), f)
				}

				total := requiredSize + framesSize + eventsSize
				if total < max {
					continue
				}

				// Reached max size, commit this range and
				// start a new range.
				startTime := time.Duration(sizes[start].Time * 1000)
				endTime := time.Duration(ev.Time * 1000)
				s.Ranges = append(s.Ranges, Range{
					Name:		fmt.Sprintf("%v-%v", startTime, endTime),
					Start:		start,
					End:		i + 1,
					StartTime:	int64(startTime),
					EndTime:	int64(endTime),
				})
				start = i + 1
				frames = make(map[string]format.Frame)
				framesSize = 0
				eventsSize = 0
			}
			if len(s.Ranges) <= 1 {
				s.Ranges = nil
				return
			}

			if end := len(sizes) - 1; start < end {
				s.Ranges = append(s.Ranges, Range{
					Name:		fmt.Sprintf("%v-%v", time.Duration(sizes[start].Time*1000), time.Duration(sizes[end].Time*1000)),
					Start:		start,
					End:		end,
					StartTime:	int64(sizes[start].Time * 1000),
					EndTime:	int64(sizes[end].Time * 1000),
				})
			}
		},
	}
}

type splitter struct {
	Ranges []Range
}

type countingWriter struct {
	size int
}

func (cw *countingWriter) Write(data []byte) (int, error) {
	cw.size += len(data)
	return len(data), nil
}

func stackFrameEncodedSize(id uint, f format.Frame) int {
	// We want to know the marginal size of traceviewer.Data.Frames for
	// each event. Running full JSON encoding of the map for each event is
	// far too slow.
	//
	// Since the format is fixed, we can easily compute the size without
	// encoding.
	//
	// A single entry looks like one of the following:
	//
	//   "1":{"name":"main.main:30"},
	//   "10":{"name":"pkg.NewSession:173","parent":9},
	//
	// The parent is omitted if 0. The trailing comma is omitted from the
	// last entry, but we don't need that much precision.
	const (
		baseSize	= len(`"`) + len(`":{"name":"`) + len(`"},`)

		// Don't count the trailing quote on the name, as that is
		// counted in baseSize.
		parentBaseSize	= len(`,"parent":`)
	)

	size := baseSize

	size += len(f.Name)

	// Bytes for id (always positive).
	for id > 0 {
		size += 1
		id /= 10
	}

	if f.Parent > 0 {
		size += parentBaseSize
		// Bytes for parent (always positive).
		for f.Parent > 0 {
			size += 1
			f.Parent /= 10
		}
	}

	return size
}

// WalkStackFrames calls fn for id and all of its parent frames from allFrames.
func WalkStackFrames(allFrames map[string]format.Frame, id int, fn func(id int)) {
	for id != 0 {
		f, ok := allFrames[strconv.Itoa(id)]
		if !ok {
			break
		}
		fn(id)
		id = f.Parent
	}
}

type Mode int

const (
	ModeGoroutineOriented	Mode	= 1 << iota
	ModeTaskOriented
	ModeThreadOriented	// Mutually exclusive with ModeGoroutineOriented.
)

// NewEmitter returns a new Emitter that writes to c. The rangeStart and
// rangeEnd args are used for splitting large traces.
func NewEmitter(c TraceConsumer, rangeStart, rangeEnd time.Duration) *Emitter {
	c.ConsumeTimeUnit("ns")

	return &Emitter{
		c:		c,
		rangeStart:	rangeStart,
		rangeEnd:	rangeEnd,
		frameTree:	frameNode{children: make(map[uint64]frameNode)},
		resources:	make(map[uint64]string),
		tasks:		make(map[uint64]task),
	}
}

type Emitter struct {
	c		TraceConsumer
	rangeStart	time.Duration
	rangeEnd	time.Duration

	heapStats, prevHeapStats	heapStats
	gstates, prevGstates		[gStateCount]int64
	threadStats, prevThreadStats	[threadStateCount]int64
	gomaxprocs			uint64
	frameTree			frameNode
	frameSeq			int
	arrowSeq			uint64
	filter				func(uint64) bool
	resourceType			string
	resources			map[uint64]string
	focusResource			uint64
	tasks				map[uint64]task
	asyncSliceSeq			uint64
}

type task struct {
	name		string
	sortIndex	int
}

func (e *Emitter) Gomaxprocs(v uint64) {
	if v > e.gomaxprocs {
		e.gomaxprocs = v
	}
}

func (e *Emitter) Resource(id uint64, name string) {
	if e.filter != nil && !e.filter(id) {
		return
	}
	e.resources[id] = name
}

func (e *Emitter) SetResourceType(name string) {
	e.resourceType = name
}

func (e *Emitter) SetResourceFilter(filter func(uint64) bool) {
	e.filter = filter
}

func (e *Emitter) Task(id uint64, name string, sortIndex int) {
	e.tasks[id] = task{name, sortIndex}
}

func (e *Emitter) Slice(s SliceEvent) {
	if e.filter != nil && !e.filter(s.Resource) {
		return
	}
	e.slice(s, format.ProcsSection, "")
}

func (e *Emitter) TaskSlice(s SliceEvent) {
	e.slice(s, format.TasksSection, pickTaskColor(s.Resource))
}

func (e *Emitter) slice(s SliceEvent, sectionID uint64, cname string) {
	if !e.tsWithinRange(s.Ts) && !e.tsWithinRange(s.Ts+s.Dur) {
		return
	}
	e.OptionalEvent(&format.Event{
		Name:		s.Name,
		Phase:		"X",
		Time:		viewerTime(s.Ts),
		Dur:		viewerTime(s.Dur),
		PID:		sectionID,
		TID:		s.Resource,
		Stack:		s.Stack,
		EndStack:	s.EndStack,
		Arg:		s.Arg,
		Cname:		cname,
	})
}

type SliceEvent struct {
	Name		string
	Ts		time.Duration
	Dur		time.Duration
	Resource	uint64
	Stack		int
	EndStack	int
	Arg		any
}

func (e *Emitter) AsyncSlice(s AsyncSliceEvent) {
	if !e.tsWithinRange(s.Ts) && !e.tsWithinRange(s.Ts+s.Dur) {
		return
	}
	if e.filter != nil && !e.filter(s.Resource) {
		return
	}
	cname := ""
	if s.TaskColorIndex != 0 {
		cname = pickTaskColor(s.TaskColorIndex)
	}
	e.asyncSliceSeq++
	e.OptionalEvent(&format.Event{
		Category:	s.Category,
		Name:		s.Name,
		Phase:		"b",
		Time:		viewerTime(s.Ts),
		TID:		s.Resource,
		ID:		e.asyncSliceSeq,
		Scope:		s.Scope,
		Stack:		s.Stack,
		Cname:		cname,
	})
	e.OptionalEvent(&format.Event{
		Category:	s.Category,
		Name:		s.Name,
		Phase:		"e",
		Time:		viewerTime(s.Ts + s.Dur),
		TID:		s.Resource,
		ID:		e.asyncSliceSeq,
		Scope:		s.Scope,
		Stack:		s.EndStack,
		Arg:		s.Arg,
		Cname:		cname,
	})
}

type AsyncSliceEvent struct {
	SliceEvent
	Category	string
	Scope		string
	TaskColorIndex	uint64	// Take on the same color as the task with this ID.
}

func (e *Emitter) Instant(i InstantEvent) {
	if !e.tsWithinRange(i.Ts) {
		return
	}
	if e.filter != nil && !e.filter(i.Resource) {
		return
	}
	cname := ""
	e.OptionalEvent(&format.Event{
		Name:		i.Name,
		Category:	i.Category,
		Phase:		"I",
		Scope:		"t",
		Time:		viewerTime(i.Ts),
		PID:		format.ProcsSection,
		TID:		i.Resource,
		Stack:		i.Stack,
		Cname:		cname,
		Arg:		i.Arg,
	})
}

type InstantEvent struct {
	Ts		time.Duration
	Name		string
	Category	string
	Resource	uint64
	Stack		int
	Arg		any
}

func (e *Emitter) Arrow(a ArrowEvent) {
	if e.filter != nil && (!e.filter(a.FromResource) || !e.filter(a.ToResource)) {
		return
	}
	e.arrow(a, format.ProcsSection)
}

func (e *Emitter) TaskArrow(a ArrowEvent) {
	e.arrow(a, format.TasksSection)
}

func (e *Emitter) arrow(a ArrowEvent, sectionID uint64) {
	if !e.tsWithinRange(a.Start) || !e.tsWithinRange(a.End) {
		return
	}
	e.arrowSeq++
	e.OptionalEvent(&format.Event{
		Name:	a.Name,
		Phase:	"s",
		TID:	a.FromResource,
		PID:	sectionID,
		ID:	e.arrowSeq,
		Time:	viewerTime(a.Start),
		Stack:	a.FromStack,
	})
	e.OptionalEvent(&format.Event{
		Name:	a.Name,
		Phase:	"t",
		TID:	a.ToResource,
		PID:	sectionID,
		ID:	e.arrowSeq,
		Time:	viewerTime(a.End),
	})
}

type ArrowEvent struct {
	Name		string
	Start		time.Duration
	End		time.Duration
	FromResource	uint64
	FromStack	int
	ToResource	uint64
}

func (e *Emitter) Event(ev *format.Event) {
	e.c.ConsumeViewerEvent(ev, true)
}

func (e *Emitter) HeapAlloc(ts time.Duration, v uint64) {
	e.heapStats.heapAlloc = v
	e.emitHeapCounters(ts)
}

func (e *Emitter) Focus(id uint64) {
	e.focusResource = id
}

func (e *Emitter) GoroutineTransition(ts time.Duration, from, to GState) {
	e.gstates[from]--
	e.gstates[to]++
	if e.prevGstates == e.gstates {
		return
	}
	if e.tsWithinRange(ts) {
		e.OptionalEvent(&format.Event{
			Name:	"Goroutines",
			Phase:	"C",
			Time:	viewerTime(ts),
			PID:	1,
			Arg: &format.GoroutineCountersArg{
				Running:	uint64(e.gstates[GRunning]),
				Runnable:	uint64(e.gstates[GRunnable]),
				GCWaiting:	uint64(e.gstates[GWaitingGC]),
			},
		})
	}
	e.prevGstates = e.gstates
}

func (e *Emitter) IncThreadStateCount(ts time.Duration, state ThreadState, delta int64) {
	e.threadStats[state] += delta
	if e.prevThreadStats == e.threadStats {
		return
	}
	if e.tsWithinRange(ts) {
		e.OptionalEvent(&format.Event{
			Name:	"Threads",
			Phase:	"C",
			Time:	viewerTime(ts),
			PID:	1,
			Arg: &format.ThreadCountersArg{
				Running:	int64(e.threadStats[ThreadStateRunning]),
				InSyscall:	int64(e.threadStats[ThreadStateInSyscall]),
				// TODO(mknyszek): Why is InSyscallRuntime not included here?
			},
		})
	}
	e.prevThreadStats = e.threadStats
}

func (e *Emitter) HeapGoal(ts time.Duration, v uint64) {
	// This cutoff at 1 PiB is a Workaround for https://github.com/golang/go/issues/63864.
	//
	// TODO(mknyszek): Remove this once the problem has been fixed.
	const PB = 1 << 50
	if v > PB {
		v = 0
	}
	e.heapStats.nextGC = v
	e.emitHeapCounters(ts)
}

func (e *Emitter) emitHeapCounters(ts time.Duration) {
	if e.prevHeapStats == e.heapStats {
		return
	}
	diff := uint64(0)
	if e.heapStats.nextGC > e.heapStats.heapAlloc {
		diff = e.heapStats.nextGC - e.heapStats.heapAlloc
	}
	if e.tsWithinRange(ts) {
		e.OptionalEvent(&format.Event{
			Name:	"Heap",
			Phase:	"C",
			Time:	viewerTime(ts),
			PID:	1,
			Arg:	&format.HeapCountersArg{Allocated: e.heapStats.heapAlloc, NextGC: diff},
		})
	}
	e.prevHeapStats = e.heapStats
}

// Err returns an error if the emitter is in an invalid state.
func (e *Emitter) Err() error {
	if e.gstates[GRunnable] < 0 || e.gstates[GRunning] < 0 || e.threadStats[ThreadStateInSyscall] < 0 || e.threadStats[ThreadStateInSyscallRuntime] < 0 {
		return fmt.Errorf(
			"runnable=%d running=%d insyscall=%d insyscallRuntime=%d",
			e.gstates[GRunnable],
			e.gstates[GRunning],
			e.threadStats[ThreadStateInSyscall],
			e.threadStats[ThreadStateInSyscallRuntime],
		)
	}
	return nil
}

func (e *Emitter) tsWithinRange(ts time.Duration) bool {
	return e.rangeStart <= ts && ts <= e.rangeEnd
}

// OptionalEvent emits ev if it's within the time range of of the consumer, i.e.
// the selected trace split range.
func (e *Emitter) OptionalEvent(ev *format.Event) {
	e.c.ConsumeViewerEvent(ev, false)
}

func (e *Emitter) Flush() {
	e.processMeta(format.StatsSection, "STATS", 0)

	if len(e.tasks) != 0 {
		e.processMeta(format.TasksSection, "TASKS", 1)
	}
	for id, task := range e.tasks {
		e.threadMeta(format.TasksSection, id, task.name, task.sortIndex)
	}

	e.processMeta(format.ProcsSection, e.resourceType, 2)

	e.threadMeta(format.ProcsSection, trace.GCP, "GC", -6)
	e.threadMeta(format.ProcsSection, trace.NetpollP, "Network", -5)
	e.threadMeta(format.ProcsSection, trace.TimerP, "Timers", -4)
	e.threadMeta(format.ProcsSection, trace.SyscallP, "Syscalls", -3)

	for id, name := range e.resources {
		priority := int(id)
		if e.focusResource != 0 && id == e.focusResource {
			// Put the focus goroutine on top.
			priority = -2
		}
		e.threadMeta(format.ProcsSection, id, name, priority)
	}

	e.c.Flush()
}

func (e *Emitter) threadMeta(sectionID, tid uint64, name string, priority int) {
	e.Event(&format.Event{
		Name:	"thread_name",
		Phase:	"M",
		PID:	sectionID,
		TID:	tid,
		Arg:	&format.NameArg{Name: name},
	})
	e.Event(&format.Event{
		Name:	"thread_sort_index",
		Phase:	"M",
		PID:	sectionID,
		TID:	tid,
		Arg:	&format.SortIndexArg{Index: priority},
	})
}

func (e *Emitter) processMeta(sectionID uint64, name string, priority int) {
	e.Event(&format.Event{
		Name:	"process_name",
		Phase:	"M",
		PID:	sectionID,
		Arg:	&format.NameArg{Name: name},
	})
	e.Event(&format.Event{
		Name:	"process_sort_index",
		Phase:	"M",
		PID:	sectionID,
		Arg:	&format.SortIndexArg{Index: priority},
	})
}

// Stack emits the given frames and returns a unique id for the stack. No
// pointers to the given data are being retained beyond the call to Stack.
func (e *Emitter) Stack(stk []*trace.Frame) int {
	return e.buildBranch(e.frameTree, stk)
}

// buildBranch builds one branch in the prefix tree rooted at ctx.frameTree.
func (e *Emitter) buildBranch(parent frameNode, stk []*trace.Frame) int {
	if len(stk) == 0 {
		return parent.id
	}
	last := len(stk) - 1
	frame := stk[last]
	stk = stk[:last]

	node, ok := parent.children[frame.PC]
	if !ok {
		e.frameSeq++
		node.id = e.frameSeq
		node.children = make(map[uint64]frameNode)
		parent.children[frame.PC] = node
		e.c.ConsumeViewerFrame(strconv.Itoa(node.id), format.Frame{Name: fmt.Sprintf("%v:%v", frame.Fn, frame.Line), Parent: parent.id})
	}
	return e.buildBranch(node, stk)
}

type heapStats struct {
	heapAlloc	uint64
	nextGC		uint64
}

func viewerTime(t time.Duration) float64 {
	return float64(t) / float64(time.Microsecond)
}

type GState int

const (
	GDead	GState	= iota
	GRunnable
	GRunning
	GWaiting
	GWaitingGC

	gStateCount
)

type ThreadState int

const (
	ThreadStateInSyscall	ThreadState	= iota
	ThreadStateInSyscallRuntime
	ThreadStateRunning

	threadStateCount
)

type frameNode struct {
	id		int
	children	map[uint64]frameNode
}

// Mapping from more reasonable color names to the reserved color names in
// https://github.com/catapult-project/catapult/blob/master/tracing/tracing/base/color_scheme.html#L50
// The chrome trace viewer allows only those as cname values.
const (
	colorLightMauve		= "thread_state_uninterruptible"	// 182, 125, 143
	colorOrange		= "thread_state_iowait"			// 255, 140, 0
	colorSeafoamGreen	= "thread_state_running"		// 126, 200, 148
	colorVistaBlue		= "thread_state_runnable"		// 133, 160, 210
	colorTan		= "thread_state_unknown"		// 199, 155, 125
	colorIrisBlue		= "background_memory_dump"		// 0, 180, 180
	colorMidnightBlue	= "light_memory_dump"			// 0, 0, 180
	colorDeepMagenta	= "detailed_memory_dump"		// 180, 0, 180
	colorBlue		= "vsync_highlight_color"		// 0, 0, 255
	colorGrey		= "generic_work"			// 125, 125, 125
	colorGreen		= "good"				// 0, 125, 0
	colorDarkGoldenrod	= "bad"					// 180, 125, 0
	colorPeach		= "terrible"				// 180, 0, 0
	colorBlack		= "black"				// 0, 0, 0
	colorLightGrey		= "grey"				// 221, 221, 221
	colorWhite		= "white"				// 255, 255, 255
	colorYellow		= "yellow"				// 255, 255, 0
	colorOlive		= "olive"				// 100, 100, 0
	colorCornflowerBlue	= "rail_response"			// 67, 135, 253
	colorSunsetOrange	= "rail_animation"			// 244, 74, 63
	colorTangerine		= "rail_idle"				// 238, 142, 0
	colorShamrockGreen	= "rail_load"				// 13, 168, 97
	colorGreenishYellow	= "startup"				// 230, 230, 0
	colorDarkGrey		= "heap_dump_stack_frame"		// 128, 128, 128
	colorTawny		= "heap_dump_child_node_arrow"		// 204, 102, 0
	colorLemon		= "cq_build_running"			// 255, 255, 119
	colorLime		= "cq_build_passed"			// 153, 238, 102
	colorPink		= "cq_build_failed"			// 238, 136, 136
	colorSilver		= "cq_build_abandoned"			// 187, 187, 187
	colorManzGreen		= "cq_build_attempt_runnig"		// 222, 222, 75
	colorKellyGreen		= "cq_build_attempt_passed"		// 108, 218, 35
	colorAnotherGrey	= "cq_build_attempt_failed"		// 187, 187, 187
)

var colorForTask = []string{
	colorLightMauve,
	colorOrange,
	colorSeafoamGreen,
	colorVistaBlue,
	colorTan,
	colorMidnightBlue,
	colorIrisBlue,
	colorDeepMagenta,
	colorGreen,
	colorDarkGoldenrod,
	colorPeach,
	colorOlive,
	colorCornflowerBlue,
	colorSunsetOrange,
	colorTangerine,
	colorShamrockGreen,
	colorTawny,
	colorLemon,
	colorLime,
	colorPink,
	colorSilver,
	colorManzGreen,
	colorKellyGreen,
}

func pickTaskColor(id uint64) string {
	idx := id % uint64(len(colorForTask))
	return colorForTask[idx]
}