1// Copyright 2023 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 counter
6
7import (
8	"fmt"
9	"runtime"
10	"strings"
11	"sync"
12)
13
14// On the disk, and upstream, stack counters look like sets of
15// regular counters with names that include newlines.
16
17// a StackCounter is the in-memory knowledge about a stack counter.
18// StackCounters are more expensive to use than regular Counters,
19// requiring, at a minimum, a call to runtime.Callers.
20type StackCounter struct {
21	name  string
22	depth int
23	file  *file
24
25	mu sync.Mutex
26	// as this is a detail of the implementation, it could be replaced
27	// by a more efficient mechanism
28	stacks []stack
29}
30
31type stack struct {
32	pcs     []uintptr
33	counter *Counter
34}
35
36func NewStack(name string, depth int) *StackCounter {
37	return &StackCounter{name: name, depth: depth, file: &defaultFile}
38}
39
40// Inc increments a stack counter. It computes the caller's stack and
41// looks up the corresponding counter. It then increments that counter,
42// creating it if necessary.
43func (c *StackCounter) Inc() {
44	pcs := make([]uintptr, c.depth)
45	n := runtime.Callers(2, pcs) // caller of Inc
46	pcs = pcs[:n]
47
48	c.mu.Lock()
49	defer c.mu.Unlock()
50
51	// Existing counter?
52	var ctr *Counter
53	for _, s := range c.stacks {
54		if eq(s.pcs, pcs) {
55			if s.counter != nil {
56				ctr = s.counter
57				break
58			}
59		}
60	}
61
62	if ctr == nil {
63		// Create new counter.
64		ctr = &Counter{
65			name: EncodeStack(pcs, c.name),
66			file: c.file,
67		}
68		c.stacks = append(c.stacks, stack{pcs: pcs, counter: ctr})
69	}
70
71	ctr.Inc()
72}
73
74// EncodeStack returns the name of the counter to
75// use for the given stack of program counters.
76// The name encodes the stack.
77func EncodeStack(pcs []uintptr, prefix string) string {
78	var locs []string
79	lastImport := ""
80	frs := runtime.CallersFrames(pcs)
81	for {
82		fr, more := frs.Next()
83		// TODO(adonovan): this CutLast(".") operation isn't
84		// appropriate for generic function symbols.
85		path, fname := cutLastDot(fr.Function)
86		if path == lastImport {
87			path = `"` // (a ditto mark)
88		} else {
89			lastImport = path
90		}
91		var loc string
92		if fr.Func != nil {
93			// Use function-relative line numbering.
94			// f:+2 means two lines into function f.
95			// f:-1 should never happen, but be conservative.
96			_, entryLine := fr.Func.FileLine(fr.Entry)
97			loc = fmt.Sprintf("%s.%s:%+d", path, fname, fr.Line-entryLine)
98		} else {
99			// The function is non-Go code or is fully inlined:
100			// use absolute line number within enclosing file.
101			loc = fmt.Sprintf("%s.%s:=%d", path, fname, fr.Line)
102		}
103		locs = append(locs, loc)
104		if !more {
105			break
106		}
107	}
108
109	name := prefix + "\n" + strings.Join(locs, "\n")
110	if len(name) > maxNameLen {
111		const bad = "\ntruncated\n"
112		name = name[:maxNameLen-len(bad)] + bad
113	}
114	return name
115}
116
117// DecodeStack expands the (compressed) stack encoded in the counter name.
118func DecodeStack(ename string) string {
119	if !strings.Contains(ename, "\n") {
120		return ename // not a stack counter
121	}
122	lines := strings.Split(ename, "\n")
123	var lastPath string // empty or ends with .
124	for i, line := range lines {
125		path, rest := cutLastDot(line)
126		if len(path) == 0 {
127			continue // unchanged
128		}
129		if len(path) == 1 && path[0] == '"' {
130			lines[i] = lastPath + rest
131		} else {
132			lastPath = path + "."
133			// line unchanged
134		}
135	}
136	return strings.Join(lines, "\n") // trailing \n?
137}
138
139// input is <import path>.<function name>
140// output is (import path, function name)
141func cutLastDot(x string) (before, after string) {
142	i := strings.LastIndex(x, ".")
143	if i < 0 {
144		return "", x
145	}
146	return x[:i], x[i+1:]
147}
148
149// Names reports all the counter names associated with a StackCounter.
150func (c *StackCounter) Names() []string {
151	c.mu.Lock()
152	defer c.mu.Unlock()
153	names := make([]string, len(c.stacks))
154	for i, s := range c.stacks {
155		names[i] = s.counter.Name()
156	}
157	return names
158}
159
160// Counters returns the known Counters for a StackCounter.
161// There may be more in the count file.
162func (c *StackCounter) Counters() []*Counter {
163	c.mu.Lock()
164	defer c.mu.Unlock()
165	counters := make([]*Counter, len(c.stacks))
166	for i, s := range c.stacks {
167		counters[i] = s.counter
168	}
169	return counters
170}
171
172func eq(a, b []uintptr) bool {
173	if len(a) != len(b) {
174		return false
175	}
176	for i := range a {
177		if a[i] != b[i] {
178			return false
179		}
180	}
181	return true
182}
183
184// ReadStack reads the given stack counter.
185// This is the implementation of
186// golang.org/x/telemetry/counter/countertest.ReadStackCounter.
187func ReadStack(c *StackCounter) (map[string]uint64, error) {
188	ret := map[string]uint64{}
189	for _, ctr := range c.Counters() {
190		v, err := Read(ctr)
191		if err != nil {
192			return nil, err
193		}
194		ret[DecodeStack(ctr.Name())] = v
195	}
196	return ret, nil
197}
198
199// IsStackCounter reports whether the counter name is for a stack counter.
200func IsStackCounter(name string) bool {
201	return strings.Contains(name, "\n")
202}
203