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

     1  // Copyright 2019 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 metrics defines a set of primitives for declaring and
     6  // managing metrics within Bigslice. Users declare metrics (such as a
     7  // counter) using the registration mechanisms provided by this
     8  // package (e.g., NewCounter). These return handles that are used for
     9  // metric operations (e.g., incrementing a counter).
    10  //
    11  // Every operation on a metric is performed in a Scope. Scopes are
    12  // provided by the Bigslice runtime and represent an operational
    13  // scope in which the metric is aggregated. For example, Bigslice
    14  // defines a Scope that is attached to each task scheduled by the
    15  // system. Scopes are merged by the Bigslice runtime to provide
    16  // aggregated metrics across larger operations (e.g., a single
    17  // session.Run).
    18  //
    19  // User functions called by Bigslice are supplied a scope through the
    20  // optional context.Context argument. The user must retrieve this
    21  // Scope using the ContextScope func.
    22  //
    23  // Metrics cannot be declared concurrently.
    24  package metrics
    25  
    26  import (
    27  	"encoding/gob"
    28  	"sync/atomic"
    29  )
    30  
    31  // metrics maps all registered metrics by id. We reserve index 0 to minimize
    32  // the chances of zero-valued metrics instances being used uninitialized.
    33  var metrics = []Metric{zeroMetric{}}
    34  
    35  // newMetric defines a new metric.
    36  func newMetric(makeMetric func(id int) Metric) {
    37  	metrics = append(metrics, makeMetric(len(metrics)))
    38  }
    39  
    40  // Metric is the abstract type of a metric. Each metric type must implement a
    41  // set of generic operations; the metric-specific operations are provided by the
    42  // metric types themselves.
    43  //
    44  // TODO(marius): eventually consider opening up this interface to allow users to
    45  // provide their own metrics implementations.
    46  type Metric interface {
    47  	// metricID is the registered ID of the metric.
    48  	metricID() int
    49  	// newInstance creates a new instance of this metric.
    50  	// Instances are managed by Scopes.
    51  	newInstance() interface{}
    52  	// merge merges the second metric instance into the first.
    53  	merge(interface{}, interface{})
    54  }
    55  
    56  // Counter is a simple counter metric. Counters implement atomic
    57  // addition and subtraction on top of an int64.
    58  type Counter struct {
    59  	id int
    60  }
    61  
    62  // NewCounter creates, registers, and returns a new Counter metric.
    63  func NewCounter() Counter {
    64  	var c Counter
    65  	newMetric(func(id int) Metric {
    66  		c.id = id
    67  		return c
    68  	})
    69  	return c
    70  }
    71  
    72  // Value retrieves the current value of this metric in the provided scope.
    73  func (c Counter) Value(scope *Scope) int64 {
    74  	return scope.instance(c).(*counterValue).load()
    75  }
    76  
    77  // Incr increments this counter's value in the provided scope by n.
    78  func (c Counter) Incr(scope *Scope, n int64) {
    79  	scope.instance(c).(*counterValue).incr(n)
    80  }
    81  
    82  // metricID implements Metric.
    83  func (c Counter) metricID() int { return c.id }
    84  
    85  // newInstance implements Metric.
    86  func (c Counter) newInstance() interface{} {
    87  	return new(counterValue)
    88  }
    89  
    90  // merge implements Metric.
    91  func (c Counter) merge(x, y interface{}) {
    92  	x.(*counterValue).merge(y.(*counterValue))
    93  }
    94  
    95  func init() {
    96  	gob.Register(&counterValue{})
    97  }
    98  
    99  // counterValue holds a single counter value. This is abstracted as its own
   100  // struct only to control how gob encodes these values. If given an
   101  // interface{}-int64, gob will happily flatten the pointers, so that they are
   102  // decoded as the wrong value (int64, not *int64). This is a workaround to avoid
   103  // this fate.
   104  type counterValue struct {
   105  	Value int64
   106  }
   107  
   108  func (c *counterValue) incr(n int64) {
   109  	atomic.AddInt64(&c.Value, n)
   110  }
   111  
   112  func (c *counterValue) load() int64 {
   113  	return atomic.LoadInt64(&c.Value)
   114  }
   115  
   116  func (c *counterValue) merge(d *counterValue) {
   117  	atomic.AddInt64(&c.Value, d.load())
   118  }
   119  
   120  // zeroMetric is used to occupy the 0th metric,
   121  // in order to help catch zero initialization bugs.
   122  type zeroMetric struct{}
   123  
   124  func (zeroMetric) metricID() int                  { return 0 }
   125  func (zeroMetric) newInstance() interface{}       { return nil }
   126  func (zeroMetric) merge(interface{}, interface{}) {}