1// Copyright 2024 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 pgo contains the compiler-agnostic portions of PGO profile handling.
6// Notably, parsing pprof profiles and serializing/deserializing from a custom
7// intermediate representation.
8package pgo
9
10import (
11	"errors"
12	"fmt"
13	"internal/profile"
14	"io"
15	"sort"
16)
17
18// FromPProf parses Profile from a pprof profile.
19func FromPProf(r io.Reader) (*Profile, error) {
20	p, err := profile.Parse(r)
21	if errors.Is(err, profile.ErrNoData) {
22		// Treat a completely empty file the same as a profile with no
23		// samples: nothing to do.
24		return emptyProfile(), nil
25	} else if err != nil {
26		return nil, fmt.Errorf("error parsing profile: %w", err)
27	}
28
29	if len(p.Sample) == 0 {
30		// We accept empty profiles, but there is nothing to do.
31		return emptyProfile(), nil
32	}
33
34	valueIndex := -1
35	for i, s := range p.SampleType {
36		// Samples count is the raw data collected, and CPU nanoseconds is just
37		// a scaled version of it, so either one we can find is fine.
38		if (s.Type == "samples" && s.Unit == "count") ||
39			(s.Type == "cpu" && s.Unit == "nanoseconds") {
40			valueIndex = i
41			break
42		}
43	}
44
45	if valueIndex == -1 {
46		return nil, fmt.Errorf(`profile does not contain a sample index with value/type "samples/count" or cpu/nanoseconds"`)
47	}
48
49	g := profile.NewGraph(p, &profile.Options{
50		SampleValue: func(v []int64) int64 { return v[valueIndex] },
51	})
52
53	namedEdgeMap, totalWeight, err := createNamedEdgeMap(g)
54	if err != nil {
55		return nil, err
56	}
57
58	if totalWeight == 0 {
59		return emptyProfile(), nil // accept but ignore profile with no samples.
60	}
61
62	return &Profile{
63		TotalWeight:  totalWeight,
64		NamedEdgeMap: namedEdgeMap,
65	}, nil
66}
67
68// createNamedEdgeMap builds a map of callsite-callee edge weights from the
69// profile-graph.
70//
71// Caller should ignore the profile if totalWeight == 0.
72func createNamedEdgeMap(g *profile.Graph) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
73	seenStartLine := false
74
75	// Process graph and build various node and edge maps which will
76	// be consumed by AST walk.
77	weight := make(map[NamedCallEdge]int64)
78	for _, n := range g.Nodes {
79		seenStartLine = seenStartLine || n.Info.StartLine != 0
80
81		canonicalName := n.Info.Name
82		// Create the key to the nodeMapKey.
83		namedEdge := NamedCallEdge{
84			CallerName:     canonicalName,
85			CallSiteOffset: n.Info.Lineno - n.Info.StartLine,
86		}
87
88		for _, e := range n.Out {
89			totalWeight += e.WeightValue()
90			namedEdge.CalleeName = e.Dest.Info.Name
91			// Create new entry or increment existing entry.
92			weight[namedEdge] += e.WeightValue()
93		}
94	}
95
96	if !seenStartLine {
97		// TODO(prattmic): If Function.start_line is missing we could
98		// fall back to using absolute line numbers, which is better
99		// than nothing.
100		return NamedEdgeMap{}, 0, fmt.Errorf("profile missing Function.start_line data (Go version of profiled application too old? Go 1.20+ automatically adds this to profiles)")
101	}
102	return postProcessNamedEdgeMap(weight, totalWeight)
103}
104
105func sortByWeight(edges []NamedCallEdge, weight map[NamedCallEdge]int64) {
106	sort.Slice(edges, func(i, j int) bool {
107		ei, ej := edges[i], edges[j]
108		if wi, wj := weight[ei], weight[ej]; wi != wj {
109			return wi > wj // want larger weight first
110		}
111		// same weight, order by name/line number
112		if ei.CallerName != ej.CallerName {
113			return ei.CallerName < ej.CallerName
114		}
115		if ei.CalleeName != ej.CalleeName {
116			return ei.CalleeName < ej.CalleeName
117		}
118		return ei.CallSiteOffset < ej.CallSiteOffset
119	})
120}
121
122func postProcessNamedEdgeMap(weight map[NamedCallEdge]int64, weightVal int64) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
123	if weightVal == 0 {
124		return NamedEdgeMap{}, 0, nil // accept but ignore profile with no samples.
125	}
126	byWeight := make([]NamedCallEdge, 0, len(weight))
127	for namedEdge := range weight {
128		byWeight = append(byWeight, namedEdge)
129	}
130	sortByWeight(byWeight, weight)
131
132	edgeMap = NamedEdgeMap{
133		Weight:   weight,
134		ByWeight: byWeight,
135	}
136
137	totalWeight = weightVal
138
139	return edgeMap, totalWeight, nil
140}
141