github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/stats/stats.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package stats provides collections of counters. Each counter
     6  // belongs to a snapshottable collection, and these collections can be
     7  // aggregated.
     8  package stats
     9  
    10  import (
    11  	"fmt"
    12  	"sort"
    13  	"strings"
    14  	"sync"
    15  	"sync/atomic"
    16  )
    17  
    18  // Values is a snapshot of the values in a collection.
    19  type Values map[string]int64
    20  
    21  // Copy returns a copy of the values v.
    22  func (v Values) Copy() Values {
    23  	w := make(Values)
    24  	for k, v := range v {
    25  		w[k] = v
    26  	}
    27  	return w
    28  }
    29  
    30  // String returns an abbreviated string with the values in this
    31  // snapshot sorted by key.
    32  func (v Values) String() string {
    33  	var keys []string
    34  	for key := range v {
    35  		keys = append(keys, key)
    36  	}
    37  	sort.Strings(keys)
    38  	for i, key := range keys {
    39  		keys[i] = fmt.Sprintf("%s:%d", key, v[key])
    40  	}
    41  	return strings.Join(keys, " ")
    42  }
    43  
    44  // A Map is a set of counters keyed by name.
    45  type Map struct {
    46  	mu     sync.Mutex
    47  	values map[string]*Int
    48  }
    49  
    50  // NewMap returns a fresh Map.
    51  func NewMap() *Map {
    52  	return &Map{
    53  		values: make(map[string]*Int),
    54  	}
    55  }
    56  
    57  // Int returns the counter with the provided name. The counter is
    58  // created if it does not already exist.
    59  func (m *Map) Int(name string) *Int {
    60  	m.mu.Lock()
    61  	v := m.values[name]
    62  	if v == nil {
    63  		v = new(Int)
    64  		m.values[name] = v
    65  	}
    66  	m.mu.Unlock()
    67  	return v
    68  }
    69  
    70  // AddAll adds all counters in the map to the provided snapshot.
    71  func (m *Map) AddAll(vals Values) {
    72  	m.mu.Lock()
    73  	for k, v := range m.values {
    74  		vals[k] += v.Get()
    75  	}
    76  	m.mu.Unlock()
    77  }
    78  
    79  // An Int is a integer counter. Ints can be atomically
    80  // incremented and set.
    81  type Int struct {
    82  	val int64
    83  }
    84  
    85  // Add increments v by delta.
    86  func (v *Int) Add(delta int64) {
    87  	if v == nil {
    88  		return
    89  	}
    90  	atomic.AddInt64(&v.val, delta)
    91  }
    92  
    93  // Set sets the counter's value to val.
    94  func (v *Int) Set(val int64) {
    95  	if v == nil {
    96  		return
    97  	}
    98  	atomic.StoreInt64(&v.val, val)
    99  }
   100  
   101  // Get returns the current value of a counter.
   102  func (v *Int) Get() int64 {
   103  	if v == nil {
   104  		return 0
   105  	}
   106  	return atomic.LoadInt64(&v.val)
   107  }