github.com/grafana/pyroscope@v1.18.0/pkg/usagestats/stats.go

package usagestats

import (
	"bytes"
	"context"
	"encoding/json"
	"expvar"
	"fmt"
	"io"
	"math"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/grafana/pyroscope/pkg/util/build"

	"github.com/cespare/xxhash/v2"
	jsoniter "github.com/json-iterator/go"
	prom "github.com/prometheus/prometheus/web/api/v1"
	"go.uber.org/atomic"
)

var (
	httpClient    = http.Client{Timeout: 5 * time.Second}
	usageStatsURL = "https://stats.grafana.org/phlare-usage-report"
	statsPrefix   = "github.com/grafana/pyroscope/"
	targetKey     = "target"
	editionKey    = "edition"

	createLock sync.RWMutex
)

// Report is the JSON object sent to the stats server
type Report struct {
	ClusterID              string    `json:"clusterID"`
	CreatedAt              time.Time `json:"createdAt"`
	Interval               time.Time `json:"interval"`
	IntervalPeriod         float64   `json:"intervalPeriod"`
	Target                 string    `json:"target"`
	prom.PrometheusVersion `json:"version"`
	Os                     string                 `json:"os"`
	Arch                   string                 `json:"arch"`
	Edition                string                 `json:"edition"`
	Metrics                map[string]interface{} `json:"metrics"`
}

// sendReport sends the report to the stats server
func sendReport(ctx context.Context, seed ClusterSeed, interval time.Time) error {
	report := buildReport(seed, interval)
	out, err := jsoniter.MarshalIndent(report, "", " ")
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodPost, usageStatsURL, bytes.NewBuffer(out))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient.Do(req.WithContext(ctx))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		data, err := io.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		return fmt.Errorf("failed to send usage stats: %s  body: %s", resp.Status, string(data))
	}
	return nil
}

// buildReport builds the report to be sent to the stats server
func buildReport(seed ClusterSeed, interval time.Time) Report {
	var (
		targetName  string
		editionName string
	)
	if target := expvar.Get(statsPrefix + targetKey); target != nil {
		if target, ok := target.(*expvar.String); ok {
			targetName = target.Value()
		}
	}
	if edition := expvar.Get(statsPrefix + editionKey); edition != nil {
		if edition, ok := edition.(*expvar.String); ok {
			editionName = edition.Value()
		}
	}

	return Report{
		ClusterID:         seed.UID,
		PrometheusVersion: build.GetVersion(),
		CreatedAt:         seed.CreatedAt,
		Interval:          interval,
		IntervalPeriod:    reportInterval.Seconds(),
		Os:                runtime.GOOS,
		Arch:              runtime.GOARCH,
		Target:            targetName,
		Edition:           editionName,
		Metrics:           buildMetrics(),
	}
}

// buildMetrics builds the metrics part of the report to be sent to the stats server
func buildMetrics() map[string]interface{} {
	result := map[string]interface{}{
		"memstats":      memstats(),
		"num_cpu":       runtime.NumCPU(),
		"num_goroutine": runtime.NumGoroutine(),
	}
	expvar.Do(func(kv expvar.KeyValue) {
		if !strings.HasPrefix(kv.Key, statsPrefix) || kv.Key == statsPrefix+targetKey || kv.Key == statsPrefix+editionKey {
			return
		}
		var value interface{}
		switch v := kv.Value.(type) {
		case *expvar.Int:
			value = v.Value()
		case *expvar.Float:
			value = v.Value()
		case *expvar.String:
			value = v.Value()
		case *Statistics:
			value = v.Value()
		case *MultiStatistics:
			value = v.Value()
		case *Counter:
			v.updateRate()
			value = v.Value()
			v.reset()
		case *MultiCounter:
			v.updateRate()
			value = v.Value()
			v.reset()
		case *WordCounter:
			value = v.Value()
		default:
			value = v.String()
		}
		result[strings.TrimPrefix(kv.Key, statsPrefix)] = value
	})
	return result
}

func memstats() interface{} {
	stats := new(runtime.MemStats)
	runtime.ReadMemStats(stats)
	return map[string]interface{}{
		"alloc":           stats.Alloc,
		"total_alloc":     stats.TotalAlloc,
		"sys":             stats.Sys,
		"heap_alloc":      stats.HeapAlloc,
		"heap_inuse":      stats.HeapInuse,
		"stack_inuse":     stats.StackInuse,
		"pause_total_ns":  stats.PauseTotalNs,
		"num_gc":          stats.NumGC,
		"gc_cpu_fraction": stats.GCCPUFraction,
	}
}

// NewFloat returns a new Float stats object.
// If a Float stats object with the same name already exists it is returned.
func NewFloat(name string) *expvar.Float {
	return createOrRetrieveExpvar(
		func() (*expvar.Float, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if f, ok := existing.(*expvar.Float); ok {
					return f, nil
				}
				return nil, fmt.Errorf("%v is set to a non-float value", name)
			}
			return nil, nil
		},
		func() *expvar.Float { // create
			return expvar.NewFloat(statsPrefix + name)
		},
	)
}

// NewInt returns a new Int stats object.
// If an Int stats object object with the same name already exists it is returned.
func NewInt(name string) *expvar.Int {
	return createOrRetrieveExpvar(
		func() (*expvar.Int, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if i, ok := existing.(*expvar.Int); ok {
					return i, nil
				}
				return nil, fmt.Errorf("%v is set to a non-int value", name)
			}
			return nil, nil
		},
		func() *expvar.Int { // create
			return expvar.NewInt(statsPrefix + name)
		},
	)
}

// NewString returns a new String stats object.
// If a String stats object with the same name already exists it is returned.
func NewString(name string) *expvar.String {
	return createOrRetrieveExpvar(
		func() (*expvar.String, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if s, ok := existing.(*expvar.String); ok {
					return s, nil
				}
				return nil, fmt.Errorf("%v is set to a non-string value", name)
			}
			return nil, nil
		},
		func() *expvar.String { // create
			return expvar.NewString(statsPrefix + name)
		},
	)
}

// Target sets the target name. This can be set multiple times.
func Target(target string) {
	NewString(targetKey).Set(target)
}

// Edition sets the edition name. This can be set multiple times.
func Edition(edition string) {
	NewString(editionKey).Set(edition)
}

type Statistics struct {
	min   *atomic.Float64
	max   *atomic.Float64
	count *atomic.Int64

	avg *atomic.Float64

	// require for stddev and stdvar
	mean  *atomic.Float64
	value *atomic.Float64
}

// NewStatistics returns a new Statistics object.
// Statistics object is thread-safe and compute statistics on the fly based on sample recorded.
// Available statistics are:
// - min
// - max
// - avg
// - count
// - stddev
// - stdvar
// If a Statistics object with the same name already exists it is returned.
func NewStatistics(name string) *Statistics {
	return createOrRetrieveExpvar(
		func() (*Statistics, error) { // check

			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if s, ok := existing.(*Statistics); ok {
					return s, nil
				}
				return nil, fmt.Errorf("%v is set to a non-Statistics value", name)
			}
			return nil, nil
		},
		func() *Statistics { // create
			s := &Statistics{
				min:   atomic.NewFloat64(math.Inf(0)),
				max:   atomic.NewFloat64(math.Inf(-1)),
				count: atomic.NewInt64(0),
				avg:   atomic.NewFloat64(0),
				mean:  atomic.NewFloat64(0),
				value: atomic.NewFloat64(0),
			}
			expvar.Publish(statsPrefix+name, s)
			return s
		},
	)
}

func (s *Statistics) String() string {
	b, _ := json.Marshal(s.Value())
	return string(b)
}

func (s *Statistics) Value() map[string]interface{} {
	stdvar := s.value.Load() / float64(s.count.Load())
	stddev := math.Sqrt(stdvar)
	min := s.min.Load()
	max := s.max.Load()
	result := map[string]interface{}{
		"avg":   s.avg.Load(),
		"count": s.count.Load(),
	}
	if !math.IsInf(min, 0) {
		result["min"] = min
	}
	if !math.IsInf(max, 0) {
		result["max"] = s.max.Load()
	}
	if !math.IsNaN(stddev) {
		result["stddev"] = stddev
	}
	if !math.IsNaN(stdvar) {
		result["stdvar"] = stdvar
	}
	return result
}

func (s *Statistics) Record(v float64) {
	for {
		min := s.min.Load()
		if min <= v {
			break
		}
		if s.min.CompareAndSwap(min, v) {
			break
		}
	}
	for {
		max := s.max.Load()
		if max >= v {
			break
		}
		if s.max.CompareAndSwap(max, v) {
			break
		}
	}
	for {
		avg := s.avg.Load()
		count := s.count.Load()
		mean := s.mean.Load()
		value := s.value.Load()

		delta := v - mean
		newCount := count + 1
		newMean := mean + (delta / float64(newCount))
		newValue := value + (delta * (v - newMean))
		newAvg := avg + ((v - avg) / float64(newCount))
		if s.avg.CompareAndSwap(avg, newAvg) && s.count.CompareAndSwap(count, newCount) && s.mean.CompareAndSwap(mean, newMean) && s.value.CompareAndSwap(value, newValue) {
			break
		}
	}
}

type MultiStatistics struct {
	m       sync.RWMutex
	values  map[string]*Statistics
	keyName string
}

// NewMultiStatistics returns a new MultiStatistics object.
// MultiStatistics object is thread-safe and computes statistics on the fly based on recorded samples.
// Available statistics are:
// - min
// - max
// - avg
// - count
// - stddev
// - stdvar
// If a MultiStatistics object with the same name already exists it is returned.
func NewMultiStatistics(name string, keyName string) *MultiStatistics {
	return createOrRetrieveExpvar(
		func() (*MultiStatistics, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if s, ok := existing.(*MultiStatistics); ok {
					return s, nil
				}
				return nil, fmt.Errorf("%v is set to a non-MultiStatistics value", name)
			}
			return nil, nil
		},
		func() *MultiStatistics { // create
			s := &MultiStatistics{
				values: map[string]*Statistics{
					"__total__": {
						min:   atomic.NewFloat64(math.Inf(0)),
						max:   atomic.NewFloat64(math.Inf(-1)),
						count: atomic.NewInt64(0),
						avg:   atomic.NewFloat64(0),
						mean:  atomic.NewFloat64(0),
						value: atomic.NewFloat64(0),
					},
				},
				keyName: keyName,
			}
			expvar.Publish(statsPrefix+name, s)
			return s
		},
	)
}

func (s *MultiStatistics) String() string {
	b, _ := json.Marshal(s.Value())
	return string(b)
}

func (s *MultiStatistics) Value() map[string]interface{} {
	s.m.RLock()
	defer s.m.RUnlock()
	var value map[string]interface{}
	valuesPerKey := make([]interface{}, 0, len(s.values))
	for k, v := range s.values {
		if k == "__total__" {
			value = v.Value()
		} else {
			valuesPerKey = append(valuesPerKey, map[string]interface{}{
				s.keyName: k,
				"data":    v.Value(),
			})
		}
	}
	value["drilldown"] = valuesPerKey
	return value
}

func (s *MultiStatistics) Record(v float64, key string) {
	keyStats := s.getOrCreateStatistics(key)
	keyStats.Record(v)
	s.values["__total__"].Record(v)
}

func (s *MultiStatistics) getOrCreateStatistics(key string) *Statistics {
	s.m.RLock()
	keyStats, ok := s.values[key]
	s.m.RUnlock()
	if ok {
		return keyStats
	}
	s.m.Lock()
	defer s.m.Unlock()
	keyStats, ok = s.values[key]
	if ok {
		return keyStats
	}
	keyStats = &Statistics{
		min:   atomic.NewFloat64(math.Inf(0)),
		max:   atomic.NewFloat64(math.Inf(-1)),
		count: atomic.NewInt64(0),
		avg:   atomic.NewFloat64(0),
		mean:  atomic.NewFloat64(0),
		value: atomic.NewFloat64(0),
	}
	s.values[key] = keyStats
	return keyStats
}

type Counter struct {
	total *atomic.Int64
	rate  *atomic.Float64

	resetTime time.Time
}

func createOrRetrieveExpvar[K any](check func() (*K, error), create func() *K) *K {
	// check if string exists holding read lock
	createLock.RLock()
	s, err := check()
	createLock.RUnlock()
	if err != nil {
		panic(err.Error())
	}
	if s != nil {
		return s
	}

	// acquire write lock and check again and create if still missing
	createLock.Lock()
	defer createLock.Unlock()
	s, err = check()
	if err != nil {
		panic(err.Error())
	}
	if s != nil {
		return s
	}

	return create()
}

// NewCounter returns a new Counter stats object.
// If a Counter stats object with the same name already exists it is returned.
func NewCounter(name string) *Counter {
	return createOrRetrieveExpvar(
		func() (*Counter, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if c, ok := existing.(*Counter); ok {
					return c, nil
				}
				return nil, fmt.Errorf("%v is set to a non-Counter value", name)
			}
			return nil, nil
		},
		func() *Counter { // create
			c := &Counter{
				total:     atomic.NewInt64(0),
				rate:      atomic.NewFloat64(0),
				resetTime: time.Now(),
			}
			expvar.Publish(statsPrefix+name, c)
			return c
		},
	)
}

func (c *Counter) updateRate() {
	total := c.total.Load()
	c.rate.Store(float64(total) / time.Since(c.resetTime).Seconds())
}

func (c *Counter) reset() {
	c.total.Store(0)
	c.rate.Store(0)
	c.resetTime = time.Now()
}

func (c *Counter) Inc(i int64) {
	c.total.Add(i)
}

func (c *Counter) String() string {
	b, _ := json.Marshal(c.Value())
	return string(b)
}

func (c *Counter) Value() map[string]interface{} {
	return map[string]interface{}{
		"total": c.total.Load(),
		"rate":  c.rate.Load(),
	}
}

type MultiCounter struct {
	m       sync.RWMutex
	values  map[string]*Counter
	keyName string
}

// NewMultiCounter returns a new MultiCounter stats object.
// If a NewMultiCounter stats object with the same name already exists it is returned.
func NewMultiCounter(name string, keyName string) *MultiCounter {
	return createOrRetrieveExpvar(
		func() (*MultiCounter, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if c, ok := existing.(*MultiCounter); ok {
					return c, nil
				}
				return nil, fmt.Errorf("%v is set to a non-MultiCounter value", name)
			}
			return nil, nil
		},
		func() *MultiCounter { // create
			c := &MultiCounter{
				values: map[string]*Counter{
					"__total__": {
						total:     atomic.NewInt64(0),
						rate:      atomic.NewFloat64(0),
						resetTime: time.Now(),
					},
				},
				keyName: keyName,
			}
			expvar.Publish(statsPrefix+name, c)
			return c
		},
	)
}

func (c *MultiCounter) updateRate() {
	c.m.RLock()
	defer c.m.RUnlock()
	for _, v := range c.values {
		v.updateRate()
	}
}

func (c *MultiCounter) reset() {
	c.m.RLock()
	defer c.m.RUnlock()
	for _, v := range c.values {
		v.reset()
	}
}

func (c *MultiCounter) Inc(i int64, keyValue string) {
	v := c.getOrCreateCounter(keyValue)
	v.Inc(i)
	c.values["__total__"].Inc(i)
}

func (c *MultiCounter) getOrCreateCounter(keyValue string) *Counter {
	c.m.RLock()
	v, ok := c.values[keyValue]
	c.m.RUnlock()
	if ok {
		return v
	}
	c.m.Lock()
	defer c.m.Unlock()
	v, ok = c.values[keyValue]
	if ok {
		return v
	}
	v = &Counter{
		total:     atomic.NewInt64(0),
		rate:      atomic.NewFloat64(0),
		resetTime: time.Now(),
	}
	c.values[keyValue] = v
	return v
}

func (c *MultiCounter) String() string {
	b, _ := json.Marshal(c.Value())
	return string(b)
}

func (c *MultiCounter) Value() map[string]interface{} {
	c.m.RLock()
	defer c.m.RUnlock()
	var value map[string]interface{}
	valuesPerKey := make([]interface{}, 0, len(c.values))
	for k, v := range c.values {
		if k == "__total__" {
			value = v.Value()
		} else {
			valuesPerKey = append(valuesPerKey, map[string]interface{}{
				c.keyName: k,
				"data":    v.Value(),
			})
		}
	}
	value["drilldown"] = valuesPerKey
	return value
}

type WordCounter struct {
	words sync.Map
	count *atomic.Int64
}

// NewWordCounter returns a new WordCounter stats object.
// The WordCounter object is thread-safe and counts the number of words recorded.
// If a WordCounter stats object with the same name already exists it is returned.
func NewWordCounter(name string) *WordCounter {
	return createOrRetrieveExpvar(
		func() (*WordCounter, error) { // check
			existing := expvar.Get(statsPrefix + name)
			if existing != nil {
				if w, ok := existing.(*WordCounter); ok {
					return w, nil
				}
				return nil, fmt.Errorf("%v is set to a non-WordCounter value", name)
			}
			return nil, nil
		},
		func() *WordCounter { // create
			c := &WordCounter{
				count: atomic.NewInt64(0),
				words: sync.Map{},
			}
			expvar.Publish(statsPrefix+name, c)
			return c
		},
	)
}

func (w *WordCounter) Add(word string) {
	if _, loaded := w.words.LoadOrStore(xxhash.Sum64String(word), struct{}{}); !loaded {
		w.count.Add(1)
	}
}

func (w *WordCounter) String() string {
	b, _ := json.Marshal(w.Value())
	return string(b)
}

func (w *WordCounter) Value() int64 {
	return w.count.Load()
}