Add field support to the sentry metrics.

Fields allow counter metrics to have multiple tabular values.

At most one field is supported at the moment.

PiperOrigin-RevId: 368767040

3 files changed, 247 insertions, 31 deletions

diff --git a/pkg/metric/metric.go b/pkg/metric/metric.go
index c9f9357de..2a2f0d611 100644
--- a/pkg/metric/metric.go
+++ b/pkg/metric/metric.go
@@ -38,7 +38,7 @@ var (
 )
 
 // Uint64Metric encapsulates a uint64 that represents some kind of metric to be
-// monitored.
+// monitored. We currently support metrics with at most one field.
 //
 // Metrics are not saved across save/restore and thus reset to zero on restore.
 //
@@ -46,6 +46,16 @@ var (
 type Uint64Metric struct {
 	// value is the actual value of the metric. It must be accessed atomically.
 	value uint64
+
+	// numFields is the number of metric fields. It is immutable once
+	// initialized.
+	numFields int
+
+	// mu protects the below fields.
+	mu sync.RWMutex `state:"nosave"`
+
+	// fields is the map of fields in the metric.
+	fields map[string]uint64
 }
 
 var (
@@ -97,8 +107,19 @@ type customUint64Metric struct {
 	// metadata describes the metric. It is immutable.
 	metadata *pb.MetricMetadata
 
-	// value returns the current value of the metric.
-	value func() uint64
+	// value returns the current value of the metric for the given set of
+	// fields. It takes a variadic number of field values as argument.
+	value func(fieldValues ...string) uint64
+}
+
+// Field contains the field name and allowed values for the metric which is
+// used in registration of the metric.
+type Field struct {
+	// name is the metric field name.
+	name string
+
+	// allowedValues is the list of allowed values for the field.
+	allowedValues []string
 }
 
 // RegisterCustomUint64Metric registers a metric with the given name.
@@ -109,7 +130,8 @@ type customUint64Metric struct {
 // Preconditions:
 // * name must be globally unique.
 // * Initialize/Disable have not been called.
-func RegisterCustomUint64Metric(name string, cumulative, sync bool, units pb.MetricMetadata_Units, description string, value func() uint64) error {
+// * value is expected to accept exactly len(fields) arguments.
+func RegisterCustomUint64Metric(name string, cumulative, sync bool, units pb.MetricMetadata_Units, description string, value func(...string) uint64, fields ...Field) error {
 	if initialized {
 		return ErrInitializationDone
 	}
@@ -129,13 +151,25 @@ func RegisterCustomUint64Metric(name string, cumulative, sync bool, units pb.Met
 		},
 		value: value,
 	}
+
+	// Metrics can exist without fields.
+	if len(fields) > 1 {
+		panic("Sentry metrics support at most one field")
+	}
+
+	for _, field := range fields {
+		allMetrics.m[name].metadata.Fields = append(allMetrics.m[name].metadata.Fields, &pb.MetricMetadata_Field{
+			FieldName:     field.name,
+			AllowedValues: field.allowedValues,
+		})
+	}
 	return nil
 }
 
-// MustRegisterCustomUint64Metric calls RegisterCustomUint64Metric and panics
-// if it returns an error.
-func MustRegisterCustomUint64Metric(name string, cumulative, sync bool, description string, value func() uint64) {
-	if err := RegisterCustomUint64Metric(name, cumulative, sync, pb.MetricMetadata_UNITS_NONE, description, value); err != nil {
+// MustRegisterCustomUint64Metric calls RegisterCustomUint64Metric for metrics
+// without fields and panics if it returns an error.
+func MustRegisterCustomUint64Metric(name string, cumulative, sync bool, description string, value func(...string) uint64, fields ...Field) {
+	if err := RegisterCustomUint64Metric(name, cumulative, sync, pb.MetricMetadata_UNITS_NONE, description, value, fields...); err != nil {
 		panic(fmt.Sprintf("Unable to register metric %q: %v", name, err))
 	}
 }
@@ -144,15 +178,24 @@ func MustRegisterCustomUint64Metric(name string, cumulative, sync bool, descript
 // name.
 //
 // Metrics must be statically defined (i.e., at init).
-func NewUint64Metric(name string, sync bool, units pb.MetricMetadata_Units, description string) (*Uint64Metric, error) {
-	var m Uint64Metric
-	return &m, RegisterCustomUint64Metric(name, true /* cumulative */, sync, units, description, m.Value)
+func NewUint64Metric(name string, sync bool, units pb.MetricMetadata_Units, description string, fields ...Field) (*Uint64Metric, error) {
+	m := Uint64Metric{
+		numFields: len(fields),
+	}
+
+	if m.numFields == 1 {
+		m.fields = make(map[string]uint64)
+		for _, fieldValue := range fields[0].allowedValues {
+			m.fields[fieldValue] = 0
+		}
+	}
+	return &m, RegisterCustomUint64Metric(name, true /* cumulative */, sync, units, description, m.Value, fields...)
 }
 
 // MustCreateNewUint64Metric calls NewUint64Metric and panics if it returns an
 // error.
-func MustCreateNewUint64Metric(name string, sync bool, description string) *Uint64Metric {
-	m, err := NewUint64Metric(name, sync, pb.MetricMetadata_UNITS_NONE, description)
+func MustCreateNewUint64Metric(name string, sync bool, description string, fields ...Field) *Uint64Metric {
+	m, err := NewUint64Metric(name, sync, pb.MetricMetadata_UNITS_NONE, description, fields...)
 	if err != nil {
 		panic(fmt.Sprintf("Unable to create metric %q: %v", name, err))
 	}
@@ -169,19 +212,56 @@ func MustCreateNewUint64NanosecondsMetric(name string, sync bool, description st
 	return m
 }
 
-// Value returns the current value of the metric.
-func (m *Uint64Metric) Value() uint64 {
-	return atomic.LoadUint64(&m.value)
+// Value returns the current value of the metric for the given set of fields.
+func (m *Uint64Metric) Value(fieldValues ...string) uint64 {
+	if m.numFields != len(fieldValues) {
+		panic(fmt.Sprintf("Number of fieldValues %d is not equal to the number of metric fields %d", len(fieldValues), m.numFields))
+	}
+
+	switch m.numFields {
+	case 0:
+		return atomic.LoadUint64(&m.value)
+	case 1:
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+
+		fieldValue := fieldValues[0]
+		if _, ok := m.fields[fieldValue]; !ok {
+			panic(fmt.Sprintf("Metric does not allow to have field value %s", fieldValue))
+		}
+		return m.fields[fieldValue]
+	default:
+		panic("Sentry metrics do not support more than one field")
+	}
 }
 
-// Increment increments the metric by 1.
-func (m *Uint64Metric) Increment() {
-	atomic.AddUint64(&m.value, 1)
+// Increment increments the metric field by 1.
+func (m *Uint64Metric) Increment(fieldValues ...string) {
+	m.IncrementBy(1, fieldValues...)
 }
 
 // IncrementBy increments the metric by v.
-func (m *Uint64Metric) IncrementBy(v uint64) {
-	atomic.AddUint64(&m.value, v)
+func (m *Uint64Metric) IncrementBy(v uint64, fieldValues ...string) {
+	if m.numFields != len(fieldValues) {
+		panic(fmt.Sprintf("Number of fieldValues %d is not equal to the number of metric fields %d", len(fieldValues), m.numFields))
+	}
+
+	switch m.numFields {
+	case 0:
+		atomic.AddUint64(&m.value, v)
+		return
+	case 1:
+		fieldValue := fieldValues[0]
+		m.mu.Lock()
+		defer m.mu.Unlock()
+
+		if _, ok := m.fields[fieldValue]; !ok {
+			panic(fmt.Sprintf("Metric does not allow to have field value %s", fieldValue))
+		}
+		m.fields[fieldValue] += v
+	default:
+		panic("Sentry metrics do not support more than one field")
+	}
 }
 
 // metricSet holds named metrics.
@@ -199,14 +279,30 @@ func makeMetricSet() metricSet {
 // Values returns a snapshot of all values in m.
 func (m *metricSet) Values() metricValues {
 	vals := make(metricValues)
+
 	for k, v := range m.m {
-		vals[k] = v.value()
+		fields := v.metadata.GetFields()
+		switch len(fields) {
+		case 0:
+			vals[k] = v.value()
+		case 1:
+			values := fields[0].GetAllowedValues()
+			fieldsMap := make(map[string]uint64)
+			for _, fieldValue := range values {
+				fieldsMap[fieldValue] = v.value(fieldValue)
+			}
+			vals[k] = fieldsMap
+		default:
+			panic(fmt.Sprintf("Unsupported number of metric fields: %d", len(fields)))
+		}
 	}
 	return vals
 }
 
-// metricValues contains a copy of the values of all metrics.
-type metricValues map[string]uint64
+// metricValues contains a copy of the values of all metrics. It is a map
+// with key as metric name and value can be either uint64 or map[string]uint64
+// to support metrics with one field.
+type metricValues map[string]interface{}
 
 var (
 	// emitMu protects metricsAtLastEmit and ensures that all emitted
@@ -233,14 +329,37 @@ func EmitMetricUpdate() {
 	snapshot := allMetrics.Values()
 
 	m := pb.MetricUpdate{}
+	// On the first call metricsAtLastEmit will be empty. Include all
+	// metrics then.
 	for k, v := range snapshot {
-		// On the first call metricsAtLastEmit will be empty. Include
-		// all metrics then.
-		if prev, ok := metricsAtLastEmit[k]; !ok || prev != v {
+		prev, ok := metricsAtLastEmit[k]
+		switch t := v.(type) {
+		case uint64:
+			// Metric exists and value did not change.
+			if ok && prev.(uint64) == t {
+				continue
+			}
+
 			m.Metrics = append(m.Metrics, &pb.MetricValue{
 				Name:  k,
-				Value: &pb.MetricValue_Uint64Value{v},
+				Value: &pb.MetricValue_Uint64Value{t},
 			})
+		case map[string]uint64:
+			for fieldValue, metricValue := range t {
+				// Emit data on the first call only if the field
+				// value has been incremented. For all other
+				// calls, emit data if the field value has been
+				// changed from the previous emit.
+				if (!ok && metricValue == 0) || (ok && prev.(map[string]uint64)[fieldValue] == metricValue) {
+					continue
+				}
+
+				m.Metrics = append(m.Metrics, &pb.MetricValue{
+					Name:        k,
+					FieldValues: []string{fieldValue},
+					Value:       &pb.MetricValue_Uint64Value{metricValue},
+				})
+			}
 		}
 	}
 
diff --git a/pkg/metric/metric.proto b/pkg/metric/metric.proto
index 3cc89047d..53c8b4b50 100644
--- a/pkg/metric/metric.proto
+++ b/pkg/metric/metric.proto
@@ -48,6 +48,15 @@ message MetricMetadata {
 
   // units is the units of the metric value.
   Units units = 6;
+
+  message Field {
+    string field_name = 1;
+    repeated string allowed_values = 2;
+  }
+
+  // fields contains the metric fields. Currently a metric can have at most
+  // one field.
+  repeated Field fields = 7;
 }
 
 // MetricRegistration contains the metadata for all metrics that will be in
@@ -66,6 +75,8 @@ message MetricValue {
   oneof value {
     uint64 uint64_value = 2;
   }
+
+  repeated string field_values = 4;
 }
 
 // MetricUpdate contains new values for multiple distinct metrics.
diff --git a/pkg/metric/metric_test.go b/pkg/metric/metric_test.go
index aefd0ea5c..c71dfd460 100644
--- a/pkg/metric/metric_test.go
+++ b/pkg/metric/metric_test.go
@@ -59,8 +59,9 @@ func reset() {
 }
 
 const (
-	fooDescription = "Foo!"
-	barDescription = "Bar Baz"
+	fooDescription     = "Foo!"
+	barDescription     = "Bar Baz"
+	counterDescription = "Counter"
 )
 
 func TestInitialize(t *testing.T) {
@@ -95,7 +96,7 @@ func TestInitialize(t *testing.T) {
 	foundBar := false
 	for _, m := range mr.Metrics {
 		if m.Type != pb.MetricMetadata_TYPE_UINT64 {
-			t.Errorf("Metadata %+v Type got %v want %v", m, m.Type, pb.MetricMetadata_TYPE_UINT64)
+			t.Errorf("Metadata %+v Type got %v want pb.MetricMetadata_TYPE_UINT64", m, m.Type)
 		}
 		if !m.Cumulative {
 			t.Errorf("Metadata %+v Cumulative got false want true", m)
@@ -256,3 +257,88 @@ func TestEmitMetricUpdate(t *testing.T) {
 		t.Errorf("%v: Value got %v want 1", m, uv.Uint64Value)
 	}
 }
+
+func TestEmitMetricUpdateWithFields(t *testing.T) {
+	defer reset()
+
+	field := Field{
+		name:          "weirdness_type",
+		allowedValues: []string{"weird1", "weird2"}}
+
+	counter, err := NewUint64Metric("/weirdness", false, pb.MetricMetadata_UNITS_NONE, counterDescription, field)
+	if err != nil {
+		t.Fatalf("NewUint64Metric got err %v want nil", err)
+	}
+
+	Initialize()
+
+	// Don't care about the registration metrics.
+	emitter.Reset()
+	EmitMetricUpdate()
+
+	// For metrics with fields, we do not emit data unless the value is
+	// incremented.
+	if len(emitter) != 0 {
+		t.Fatalf("EmitMetricUpdate emitted %d events want 0", len(emitter))
+	}
+
+	counter.IncrementBy(4, "weird1")
+	counter.Increment("weird2")
+
+	emitter.Reset()
+	EmitMetricUpdate()
+
+	if len(emitter) != 1 {
+		t.Fatalf("EmitMetricUpdate emitted %d events want 1", len(emitter))
+	}
+
+	update, ok := emitter[0].(*pb.MetricUpdate)
+	if !ok {
+		t.Fatalf("emitter %v got %T want pb.MetricUpdate", emitter[0], emitter[0])
+	}
+
+	if len(update.Metrics) != 2 {
+		t.Errorf("MetricUpdate got %d metrics want 2", len(update.Metrics))
+	}
+
+	foundWeird1 := false
+	foundWeird2 := false
+	for i := 0; i < len(update.Metrics); i++ {
+		m := update.Metrics[i]
+
+		if m.Name != "/weirdness" {
+			t.Errorf("Metric %+v name got %q want '/weirdness'", m, m.Name)
+		}
+		if len(m.FieldValues) != 1 {
+			t.Errorf("MetricUpdate got %d fields want 1", len(m.FieldValues))
+		}
+
+		switch m.FieldValues[0] {
+		case "weird1":
+			uv, ok := m.Value.(*pb.MetricValue_Uint64Value)
+			if !ok {
+				t.Errorf("%+v: value %v got %T want pb.MetricValue_Uint64Value", m, m.Value, m.Value)
+			}
+			if uv.Uint64Value != 4 {
+				t.Errorf("%v: Value got %v want 4", m, uv.Uint64Value)
+			}
+			foundWeird1 = true
+		case "weird2":
+			uv, ok := m.Value.(*pb.MetricValue_Uint64Value)
+			if !ok {
+				t.Errorf("%+v: value %v got %T want pb.MetricValue_Uint64Value", m, m.Value, m.Value)
+			}
+			if uv.Uint64Value != 1 {
+				t.Errorf("%v: Value got %v want 1", m, uv.Uint64Value)
+			}
+			foundWeird2 = true
+		}
+	}
+
+	if !foundWeird1 {
+		t.Errorf("Field value weird1 not found: %+v", emitter)
+	}
+	if !foundWeird2 {
+		t.Errorf("Field value weird2 not found: %+v", emitter)
+	}
+}