stackcounter.go

     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  
     5  package counter
     6  
     7  import (
     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.
    20  type 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  
    31  type stack struct {
    32  	pcs     []uintptr
    33  	counter *Counter
    34  }
    35  
    36  func 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.
    43  func (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.
    77  func 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.
   118  func 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)
   141  func 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.
   150  func (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.
   162  func (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  
   172  func 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.
   187  func 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.
   200  func IsStackCounter(name string) bool {
   201  	return strings.Contains(name, "\n")
   202  }
   203
View as plain text