From d04a8d3460146c3fd255ec130af39f05c1599bc5 Mon Sep 17 00:00:00 2001 From: Etienne Perot Date: Tue, 27 Dec 2022 15:02:47 -0800 Subject: [PATCH] gVisor: Add library for exporting instrumentation data in Prometheus format. This adds a new library, `//pkg/prometheus`, which contains just enough data structures such that we can encode instrumentation information in Prometheus information. These data structures are JSON-encodable, such that they can be used over the `runsc` control channel for export (implemented in a future CL). The existing `metric.go` library gains new functionality to export its own data using this new export format. This change is part of a series of changes to support Prometheus-style metrics in `runsc`. Doing so requires making several seemingly-odd design decisions, due to the following architectural constraints: - Prometheus requires an HTTP server serving the `/metrics` endpoint. - For performance reasons, the `runsc boot` process cannot run the `netpoller` goroutine. - Since we don't want to write our own HTTP server implementation, this means the HTTP endpoint has to be served by a separate process that remains running during the lifetime of the container. - The `runsc boot` process is untrusted. - This means we cannot trust metrics data that comes out of the Sentry. Therefore, there needs to be an elaborate dance where we pre-register metric metadata before starting any untrusted workload. Then, the server relaying the metric data must verify the validity of metric values against this metric metadata. This avoids leaking metrics, cardinality blow-ups, and other such DoS vectors. - This feature needs to be easy-to-use in a typical Docker setting. - This means having the ability to just say `--metrics-server=localhost:1337` in the `runsc` runtime entry in `/etc/docker/daemon.json` and have that Just Work(TM), even when multiple containers are running. - Since only one process may listen on a port at a given time, this means the metric server needs to be able to multiplex requests out to multiple running sandboxes, and remain alive for the entire duration of either of these sandboxes. However, it should also die when there are no sandboxes, so that we don't end up with leftover metric servers lying around. - For this reason, the metrics server runs *outside* of the usual per-container cgroups. - This also saves system resources by not running one server per sandbox. - The metrics server must be exposed to the outside world, and cannot assume that its clients are trustworthy. - For this reason, a metrics server is bound to a runtime root directory, and double-checks all that the sandboxes it is asked to follow actually exist in this root directory. PiperOrigin-RevId: 498039624 --- pkg/metric/BUILD | 1 + pkg/metric/metric.go | 120 +++++- pkg/metric/metric.proto | 4 + pkg/metric/metric_test.go | 11 +- pkg/prometheus/BUILD | 27 ++ pkg/prometheus/prometheus.go | 444 ++++++++++++++++++++ pkg/prometheus/prometheus_test.go | 653 ++++++++++++++++++++++++++++++ 7 files changed, 1247 insertions(+), 13 deletions(-) create mode 100644 pkg/prometheus/BUILD create mode 100644 pkg/prometheus/prometheus.go create mode 100644 pkg/prometheus/prometheus_test.go diff --git a/pkg/metric/BUILD b/pkg/metric/BUILD index 695def5f6..9b7eb0d98 100644 --- a/pkg/metric/BUILD +++ b/pkg/metric/BUILD @@ -18,6 +18,7 @@ go_library( "//pkg/eventchannel", "//pkg/gohacks", "//pkg/log", + "//pkg/prometheus", "//pkg/sync", "@org_golang_google_protobuf//types/known/timestamppb", ], diff --git a/pkg/metric/metric.go b/pkg/metric/metric.go index a126a8036..cc8c9689b 100644 --- a/pkg/metric/metric.go +++ b/pkg/metric/metric.go @@ -20,6 +20,7 @@ import ( "fmt" "math" "sort" + "strings" "time" "google.golang.org/protobuf/types/known/timestamppb" @@ -27,6 +28,7 @@ import ( "gvisor.dev/gvisor/pkg/eventchannel" "gvisor.dev/gvisor/pkg/log" pb "gvisor.dev/gvisor/pkg/metric/metric_go_proto" + "gvisor.dev/gvisor/pkg/prometheus" "gvisor.dev/gvisor/pkg/sync" ) @@ -165,6 +167,9 @@ type customUint64Metric struct { // metadata describes the metric. It is immutable. metadata *pb.MetricMetadata + // prometheusMetric describes the metric in Prometheus format. It is immutable. + prometheusMetric *prometheus.Metric + // 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 @@ -327,6 +332,12 @@ func (m fieldMapper) keyToMultiField(key int) []string { return fields } +// nameToPrometheusName transforms a path-style metric name (/foo/bar) into a Prometheus-style +// metric name (foo_bar). +func nameToPrometheusName(name string) string { + return strings.ReplaceAll(strings.TrimPrefix(name, "/"), "/", "_") +} + // RegisterCustomUint64Metric registers a metric with the given name. // // Register must only be called at init and will return and error if called @@ -348,14 +359,25 @@ func RegisterCustomUint64Metric(name string, cumulative, sync bool, units pb.Met return ErrNameInUse } + promType := prometheus.TypeGauge + if cumulative { + promType = prometheus.TypeCounter + } + allMetrics.uint64Metrics[name] = customUint64Metric{ metadata: &pb.MetricMetadata{ - Name: name, - Description: description, - Cumulative: cumulative, - Sync: sync, - Type: pb.MetricMetadata_TYPE_UINT64, - Units: units, + Name: name, + PrometheusName: nameToPrometheusName(name), + Description: description, + Cumulative: cumulative, + Sync: sync, + Type: pb.MetricMetadata_TYPE_UINT64, + Units: units, + }, + prometheusMetric: &prometheus.Metric{ + Name: nameToPrometheusName(name), + Help: description, + Type: promType, }, value: value, } @@ -600,15 +622,18 @@ type DistributionMetric struct { // and we call whichever one is in use in AddSample. exponentialBucketer *ExponentialBucketer - // metadata is the metadata about this metric. + // metadata is the metadata about this metric. It is immutable. metadata *pb.MetricMetadata + // prometheusMetric describes the metric in Prometheus format. It is immutable. + prometheusMetric *prometheus.Metric + // fieldsToKey converts a multi-dimensional fields to a single string to use // as key for `samples`. fieldsToKey fieldMapper // samples is the number of samples that fell within each bucket. - // It is mapped by the concatenation of the fields, using fieldsToKey. + // It is mapped by the concatenation of the fields using `fieldsToKey`. // The value is a list of bucket sample counts, with the 0-th being the // "underflow bucket", i.e. the bucket of samples which cannot fall into // any bucket that the bucketer supports. @@ -617,6 +642,10 @@ type DistributionMetric struct { // The last value is the number of samples that fell into the bucketer's // last (i.e. infinite) bucket. samples [][]atomicbitops.Uint64 + + // sampleSum is the sum of samples. + // It is mapped by the concatenation of the fields using `fieldsToKey`. + sampleSum []atomicbitops.Int64 } // NewDistributionMetric creates and registers a new distribution metric. @@ -656,8 +685,10 @@ func NewDistributionMetric(name string, sync bool, bucketer Bucketer, unit pb.Me exponentialBucketer: exponentialBucketer, fieldsToKey: fieldsToKey, samples: samples, + sampleSum: make([]atomicbitops.Int64, fieldsToKey.numKeys()), metadata: &pb.MetricMetadata{ Name: name, + PrometheusName: nameToPrometheusName(name), Description: description, Cumulative: false, Sync: sync, @@ -666,6 +697,11 @@ func NewDistributionMetric(name string, sync bool, bucketer Bucketer, unit pb.Me Fields: protoFields, DistributionBucketLowerBounds: lowerBounds, }, + prometheusMetric: &prometheus.Metric{ + Name: nameToPrometheusName(name), + Type: prometheus.TypeHistogram, + Help: description, + }, } return allMetrics.distributionMetrics[name], nil } @@ -696,6 +732,7 @@ func (d *DistributionMetric) AddSample(sample int64, fields ...string) { func (d *DistributionMetric) addSampleByKey(sample int64, key int) { bucket := d.exponentialBucketer.BucketIndex(sample) d.samples[key][bucket+1].Add(1) + d.sampleSum[key].Add(sample) } // Minimum number of buckets for NewDurationBucket. @@ -845,6 +882,7 @@ func (m *metricSet) Values() metricValues { uint64Metrics: make(map[string]any, len(m.uint64Metrics)), distributionMetrics: make(map[string][][]uint64, len(m.distributionMetrics)), distributionTotalSamples: make(map[string][]uint64, len(m.distributionMetrics)), + distributionSampleSum: make(map[string][]int64, len(m.distributionMetrics)), stages: stages, } for k, v := range m.uint64Metrics { @@ -866,6 +904,7 @@ func (m *metricSet) Values() metricValues { for name, metric := range m.distributionMetrics { fieldKeysToValues := make([][]uint64, len(metric.samples)) fieldKeysToTotalSamples := make([]uint64, len(metric.samples)) + fieldKeysToSampleSum := make([]int64, len(metric.samples)) for fieldKey, samples := range metric.samples { samplesSnapshot := snapshotDistribution(samples) totalSamples := uint64(0) @@ -877,14 +916,17 @@ func (m *metricSet) Values() metricValues { // the maps for this fieldKey as nil. This lessens the memory cost // of distributions with unused field combinations. fieldKeysToTotalSamples[fieldKey] = 0 + fieldKeysToSampleSum[fieldKey] = 0 fieldKeysToValues[fieldKey] = nil } else { fieldKeysToTotalSamples[fieldKey] = totalSamples + fieldKeysToSampleSum[fieldKey] = metric.sampleSum[fieldKey].Load() fieldKeysToValues[fieldKey] = samplesSnapshot } } vals.distributionMetrics[name] = fieldKeysToValues vals.distributionTotalSamples[name] = fieldKeysToTotalSamples + vals.distributionSampleSum[name] = fieldKeysToSampleSum } return vals } @@ -912,6 +954,10 @@ type metricValues struct { // no new samples are not retransmitted. distributionTotalSamples map[string][]uint64 + // distributionSampleSum is the sum of samples for each distribution metric + // and field values. + distributionSampleSum map[string][]int64 + // Information on when initialization stages were reached. Does not include // the currently-ongoing stage, if any. stages []stageTiming @@ -1052,6 +1098,64 @@ func EmitMetricUpdate() { } } +// GetSnapshot returns a Prometheus snapshot of the metric data. +func GetSnapshot() *prometheus.Snapshot { + values := allMetrics.Values() + snapshot := prometheus.NewSnapshot() + for k, v := range values.uint64Metrics { + m := allMetrics.uint64Metrics[k] + switch t := v.(type) { + case uint64: + snapshot.Add(prometheus.NewIntData(m.prometheusMetric, int64(t))) + case map[string]uint64: + for fieldValue, metricValue := range t { + snapshot.Add(prometheus.LabeledIntData(m.prometheusMetric, map[string]string{ + // uint64 metrics currently only support at most one field name. + m.metadata.Fields[0].GetFieldName(): fieldValue, + }, int64(metricValue))) + } + } + } + for k, dists := range values.distributionTotalSamples { + m := allMetrics.distributionMetrics[k] + distributionSamples := values.distributionMetrics[k] + numFiniteBuckets := m.exponentialBucketer.NumFiniteBuckets() + sampleSums := values.distributionSampleSum[k] + for fieldKey := range dists { + var labels map[string]string + if numFields := m.fieldsToKey.numKeys(); numFields > 0 { + labels = make(map[string]string, numFields) + for fieldIndex, field := range m.fieldsToKey.keyToMultiField(fieldKey) { + labels[m.metadata.Fields[fieldIndex].GetFieldName()] = field + } + } + currentSamples := distributionSamples[fieldKey] + buckets := make([]prometheus.Bucket, numFiniteBuckets+2) + for b := 0; b < numFiniteBuckets+2; b++ { + var upperBound prometheus.Number + if b == numFiniteBuckets+1 { + upperBound = prometheus.Number{Float: math.Inf(1)} // Overflow bucket. + } else { + upperBound = prometheus.Number{Int: m.exponentialBucketer.LowerBound(b + 1)} + } + buckets[b] = prometheus.Bucket{ + Samples: currentSamples[b], + UpperBound: upperBound, + } + } + snapshot.Add(&prometheus.Data{ + Metric: m.prometheusMetric, + Labels: labels, + HistogramValue: &prometheus.Histogram{ + Total: prometheus.Number{Int: sampleSums[fieldKey]}, + Buckets: buckets, + }, + }) + } + } + return snapshot +} + // StartStage should be called when an initialization stage is started. // It returns a function that must be called to indicate that the stage ended. // Alternatively, future calls to StartStage will implicitly indicate that the diff --git a/pkg/metric/metric.proto b/pkg/metric/metric.proto index 59dfdf6fe..0077a5aa3 100644 --- a/pkg/metric/metric.proto +++ b/pkg/metric/metric.proto @@ -24,6 +24,10 @@ message MetricMetadata { // (e.g., /foo/count). string name = 1; + // prometheus_name is the unique name of the metric in Prometheus format + // (e.g. foo_count). + string prometheus_name = 9; + // description is a human-readable description of the metric. string description = 2; diff --git a/pkg/metric/metric_test.go b/pkg/metric/metric_test.go index 4856c8259..109d6c9ce 100644 --- a/pkg/metric/metric_test.go +++ b/pkg/metric/metric_test.go @@ -113,11 +113,12 @@ func TestInitialize(t *testing.T) { case "/distrib": foundDistrib = true want := &pb.MetricMetadata{ - Name: "/distrib", - Type: pb.MetricMetadata_TYPE_DISTRIBUTION, - Units: pb.MetricMetadata_UNITS_NANOSECONDS, - Description: distribDescription, - Sync: true, + Name: "/distrib", + PrometheusName: "distrib", + Type: pb.MetricMetadata_TYPE_DISTRIBUTION, + Units: pb.MetricMetadata_UNITS_NANOSECONDS, + Description: distribDescription, + Sync: true, Fields: []*pb.MetricMetadata_Field{ {FieldName: "field1", AllowedValues: []string{"foo", "bar"}}, {FieldName: "field2", AllowedValues: []string{"baz", "quux"}}, diff --git a/pkg/prometheus/BUILD b/pkg/prometheus/BUILD new file mode 100644 index 000000000..5a2734784 --- /dev/null +++ b/pkg/prometheus/BUILD @@ -0,0 +1,27 @@ +load("//tools:defs.bzl", "go_library", "go_test") + +package(licenses = ["notice"]) + +go_library( + name = "prometheus", + srcs = [ + "prometheus.go", + ], + visibility = ["//:sandbox"], +) + +go_test( + name = "prometheus_test", + srcs = ["prometheus_test.go"], + library = ":prometheus", + deps = [ + "//pkg/metric:metric_go_proto", + "@com_github_golang_protobuf//proto:go_default_library", + "@com_github_google_go_cmp//cmp:go_default_library", + "@com_github_prometheus_common//expfmt", + "@org_golang_google_protobuf//encoding/prototext:go_default_library", + "@org_golang_google_protobuf//proto:go_default_library", + "@org_golang_google_protobuf//reflect/protoreflect:go_default_library", + "@org_golang_google_protobuf//testing/protocmp:go_default_library", + ], +) diff --git a/pkg/prometheus/prometheus.go b/pkg/prometheus/prometheus.go new file mode 100644 index 000000000..e6713483a --- /dev/null +++ b/pkg/prometheus/prometheus.go @@ -0,0 +1,444 @@ +// Copyright 2022 The gVisor Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package prometheus contains Prometheus-compliant metric data structures and utilities. +// It can export data in Prometheus data format, documented at: +// https://prometheus.io/docs/instrumenting/exposition_formats/ +package prometheus + +import ( + "bufio" + "fmt" + "io" + "math" + "sort" + "strings" + "time" +) + +// timeNow is the time.Now() function. Can be mocked in tests. +var timeNow = time.Now + +// Type is a Prometheus metric type. +type Type int + +// List of supported Prometheus metric types. +const ( + TypeUntyped = Type(iota) + TypeGauge + TypeCounter + TypeHistogram +) + +// Metric is a Prometheus metric metadata. +type Metric struct { + // Name is the Prometheus metric name. + Name string `json:"name"` + + // Type is the type of the metric. + Type Type `json:"type"` + + // Help is an optional helpful string explaining what the metric is about. + Help string `json:"help"` +} + +// writeHeaderTo writes the metric comment header to the given writer. +func (m *Metric) writeHeaderTo(w io.Writer, options ExportOptions) error { + if m.Help != "" { + // Prometheus metric description escape rules: Only backslashes and line breaks need escaping. + if _, err := io.WriteString(w, fmt.Sprintf("# HELP %s%s %s\n", options.ExporterPrefix, m.Name, strings.ReplaceAll(strings.ReplaceAll(m.Help, "\\", "\\\\"), "\n", "\\n"))); err != nil { + return err + } + } + var metricType string + switch m.Type { + case TypeGauge: + metricType = "gauge" + case TypeCounter: + metricType = "counter" + case TypeHistogram: + metricType = "histogram" + case TypeUntyped: + metricType = "untyped" + } + if metricType != "" { + if _, err := io.WriteString(w, fmt.Sprintf("# TYPE %s%s %s\n", options.ExporterPrefix, m.Name, metricType)); err != nil { + return err + } + } + return nil +} + +// Number represents a numerical value. +// In Prometheus, all numbers are float64s. +// However, for the purpose of usage of this library, we support expressing numbers as integers, +// which makes things like counters much easier and more precise. +// At data export time (i.e. when written out in Prometheus data format), it is coallesced into +// a float. +type Number struct { + // Float is the float value of this number. + // Mutually exclusive with Int. + Float float64 `json:"float,omitempty"` + + // Int is the integer value of this number. + // Mutually exclusive with Float. + Int int64 `json:"int,omitempty"` +} + +// String returns a string representation of this number. +func (n *Number) String() string { + var s strings.Builder + if err := n.writeTo(&s); err != nil { + panic(err) + } + return s.String() +} + +// writeTo writes the number to the given writer. +func (n *Number) writeTo(w io.Writer) error { + var s string + switch { + // Zero case: + case n.Int == 0 && n.Float == 0: + s = "0" + + // Integer case: + case n.Int != 0: + s = fmt.Sprintf("%d", n.Int) + + // Special float cases: + case n.Float == math.Inf(-1): + s = "-Inf" + case n.Float == math.Inf(1): + s = "+Inf" + case math.IsNaN(n.Float): + s = "NaN" + + // Regular float case: + default: + s = fmt.Sprintf("%f", n.Float) + } + _, err := io.WriteString(w, s) + return err +} + +// Bucket is a single histogram bucket. +type Bucket struct { + // UpperBound is the upper bound of the bucket. + // The lower bound of the bucket is the largest UpperBound within other Histogram Buckets that + // is smaller than this bucket's UpperBound. + // The bucket with the smallest UpperBound within a Histogram implicitly has -Inf as lower bound. + // This should be set to +Inf to mark the "last" bucket. + UpperBound Number `json:"le"` + + // Samples is the number of samples in the bucket. + // Note: When exported to Prometheus, they are exported cumulatively, i.e. the count of samples + // exported in Bucket i is actually sum(histogram.Buckets[j].Samples for 0 <= j <= i). + Samples uint64 `json:"n,omitempty"` +} + +// Histogram contains data about histogram values. +type Histogram struct { + // Total is the sum of sample values across all buckets. + Total Number `json:"total"` + // Buckets contains per-bucket data. + // A distribution with n finite-boundary buckets should have n+2 entries here. + // The 0th entry is the underflow bucket (i.e. the one with -inf as lower bound), + // and the last aka (n+1)th entry is the overflow bucket (i.e. the one with +inf as upper bound). + Buckets []Bucket `json:"buckets,omitempty"` +} + +// Data is an observation of the value of a single metric at a certain point in time. +type Data struct { + // Metric is the metric for which the value is being reported. + Metric *Metric `json:"metric"` + + // Labels is a key-value pair representing the labels set on this metric. + // This may be merged with other labels during export. + Labels map[string]string `json:"labels,omitempty"` + + // At most one of the fields below may be set. + // Which one depends on the type of the metric. + + // Number is used for all numerical types. + Number *Number `json:"val,omitempty"` + + // Histogram is used for histogram-typed metrics. + HistogramValue *Histogram `json:"histogram,omitempty"` +} + +// NewIntData returns a new Data struct with the given metric and value. +func NewIntData(metric *Metric, val int64) *Data { + return &Data{Metric: metric, Number: &Number{Int: val}} +} + +// LabeledIntData returns a new Data struct with the given metric, labels, and value. +func LabeledIntData(metric *Metric, labels map[string]string, val int64) *Data { + return &Data{Metric: metric, Labels: labels, Number: &Number{Int: val}} +} + +// NewFloatData returns a new Data struct with the given metric and value. +func NewFloatData(metric *Metric, val float64) *Data { + return &Data{Metric: metric, Number: &Number{Float: val}} +} + +// ExportOptions contains options that control how metric data is exported in Prometheus format. +type ExportOptions struct { + // CommentHeader is prepended as a comment before any metric data is exported. + CommentHeader string + + // ExporterPrefix is prepended to all metric names. + ExporterPrefix string + + // ExtraLabels is added as labels for all metric values. + ExtraLabels map[string]string +} + +// writeMetricPreambleTo writes the metric name to w. It may also write the help and type comments +// of the metric, if they haven't been written to w yet, as tracked by metricsWritten. +func (d *Data) writeMetricPreambleTo(w io.Writer, options ExportOptions, metricsWritten map[string]bool) error { + // Metric header, if we haven't printed it yet. + if !metricsWritten[d.Metric.Name] { + if err := d.Metric.writeHeaderTo(w, options); err != nil { + return err + } + metricsWritten[d.Metric.Name] = true + } + + // Metric name. + if options.ExporterPrefix != "" { + if _, err := io.WriteString(w, options.ExporterPrefix); err != nil { + return err + } + } + if _, err := io.WriteString(w, d.Metric.Name); err != nil { + return err + } + return nil +} + +// OrderedLabels returns the list of 'label_key="label_value"' in sorted order, except "le" which is +// a reserved Prometheus label name and should go last. +func OrderedLabels(labels ...map[string]string) ([]string, error) { + var le string + totalLabels := 0 + for _, labelMap := range labels { + if leVal, found := labelMap["le"]; found { + le = leVal + totalLabels += len(labelMap) - 1 + } else { + totalLabels += len(labelMap) + } + } + if le != "" { + totalLabels++ + } + keys := make(map[string]struct{}, totalLabels) + for _, labelMap := range labels { + for label := range labelMap { + if _, found := keys[label]; found { + return nil, fmt.Errorf("duplicate label name %q", label) + } + keys[label] = struct{}{} + } + } + orderedKeys := make([]string, 0, totalLabels) + for _, labelMap := range labels { + for k, v := range labelMap { + if k != "le" { + orderedKeys = append(orderedKeys, fmt.Sprintf("%s=%q", k, v)) + } + } + } + sort.Strings(orderedKeys) + if le != "" { + orderedKeys = append(orderedKeys, fmt.Sprintf("le=%q", le)) + } + return orderedKeys, nil +} + +// writeLabelsTo writes a set of metric labels. +func (d *Data) writeLabelsTo(w io.Writer, extraLabels map[string]string, leLabel *Number) error { + if (d.Labels != nil && len(d.Labels) != 0) || (extraLabels != nil && len(extraLabels) != 0) || leLabel != nil { + if _, err := io.WriteString(w, "{"); err != nil { + return err + } + var orderedLabels []string + var err error + if leLabel != nil { + orderedLabels, err = OrderedLabels(d.Labels, extraLabels, map[string]string{"le": leLabel.String()}) + } else { + orderedLabels, err = OrderedLabels(d.Labels, extraLabels) + } + if err != nil { + return err + } + for i, keyVal := range orderedLabels { + if i != 0 { + if _, err := io.WriteString(w, ","); err != nil { + return err + } + } + if _, err := io.WriteString(w, keyVal); err != nil { + return err + } + } + if _, err := io.WriteString(w, "}"); err != nil { + return err + } + } + return nil +} + +// writeMetricLine writes a single line with a single number (val) to w. +func (d *Data) writeMetricLine(w io.Writer, metricSuffix string, val *Number, when time.Time, options ExportOptions, leLabel *Number, metricsWritten map[string]bool) error { + if err := d.writeMetricPreambleTo(w, options, metricsWritten); err != nil { + return err + } + if metricSuffix != "" { + if _, err := io.WriteString(w, metricSuffix); err != nil { + return err + } + } + if err := d.writeLabelsTo(w, options.ExtraLabels, leLabel); err != nil { + return err + } + if _, err := io.WriteString(w, " "); err != nil { + return err + } + if err := val.writeTo(w); err != nil { + return err + } + if _, err := io.WriteString(w, fmt.Sprintf(" %d\n", when.UnixMilli())); err != nil { + return err + } + return nil +} + +// writeTo writes the Data to the given writer, in Prometheus format. +func (d *Data) writeTo(w io.Writer, when time.Time, options ExportOptions, metricsWritten map[string]bool) error { + switch d.Metric.Type { + case TypeUntyped, TypeGauge, TypeCounter: + return d.writeMetricLine(w, "", d.Number, when, options, nil, metricsWritten) + case TypeHistogram: + // Write an empty line before and after histograms, to make them easier to distinguish from + // other metric lines. + if _, err := io.WriteString(w, "\n"); err != nil { + return err + } + var numSamples uint64 + var samples Number + for _, bucket := range d.HistogramValue.Buckets { + numSamples += bucket.Samples + samples.Int = int64(numSamples) // Prometheus distribution bucket counts are cumulative. + if err := d.writeMetricLine(w, "_bucket", &samples, when, options, &bucket.UpperBound, metricsWritten); err != nil { + return err + } + } + if err := d.writeMetricLine(w, "_sum", &d.HistogramValue.Total, when, options, nil, metricsWritten); err != nil { + return err + } + samples.Int = int64(numSamples) + if err := d.writeMetricLine(w, "_count", &samples, when, options, nil, metricsWritten); err != nil { + return err + } + // Empty line after the histogram. + if _, err := io.WriteString(w, "\n"); err != nil { + return err + } + return nil + default: + return fmt.Errorf("unknown metric type for metric %s: %v", d.Metric.Name, d.Metric.Type) + } +} + +// Snapshot is a snapshot of the values of all the metrics at a certain point in time. +type Snapshot struct { + // When is the timestamp at which the snapshot was taken. + // Note that Prometheus ultimately encodes timestamps as millisecond-precision int64s from epoch. + When time.Time `json:"when,omitempty"` + + // Data is the whole snapshot data. + // Each Data must be a unique combination of (Metric, Labels) within a Snapshot. + Data []*Data `json:"data,omitempty"` +} + +// NewSnapshot returns a new Snapshot at the current time. +func NewSnapshot() *Snapshot { + return &Snapshot{When: timeNow()} +} + +// Add data point(s) to the snapshot. +// Returns itself for chainability. +func (s *Snapshot) Add(data ...*Data) *Snapshot { + s.Data = append(s.Data, data...) + return s +} + +// countingWriter implements io.Writer, and counts the number of bytes written to it. +// Useful in this file to keep track of total number of bytes without having to plumb this +// everywhere in the writeX() functions in this file. +type countingWriter struct { + w *bufio.Writer + written int +} + +// Write implements io.Writer.Write. +func (w *countingWriter) Write(b []byte) (int, error) { + written, err := w.w.Write(b) + w.written += written + return written, err +} + +// Written returns the number of bytes written to the underlying writer (minus buffered writes). +func (w *countingWriter) Written() int { + return w.written - w.w.Buffered() +} + +// WriteTo writes the data to the given writer, in Prometheus format. +// It returns the number of bytes written. +func (s *Snapshot) WriteTo(w io.Writer, options ExportOptions) (int, error) { + // Add wrapping to the buffer. Note that bufio is smart enough to not wrap buffered writers within + // buffered writers, so if the caller passes in a buffered writer, this won't double-buffer. + cw := &countingWriter{w: bufio.NewWriter(w)} + if options.CommentHeader != "" { + for _, commentLine := range strings.Split(options.CommentHeader, "\n") { + if _, err := io.WriteString(cw, "# "); err != nil { + return cw.Written(), err + } + if _, err := io.WriteString(cw, commentLine); err != nil { + return cw.Written(), err + } + if _, err := io.WriteString(cw, "\n"); err != nil { + return cw.Written(), err + } + } + } + if _, err := io.WriteString(cw, fmt.Sprintf("# Metric snapshot containing %d data points taken at: %v\n", len(s.Data), s.When)); err != nil { + return cw.Written(), err + } + metricsWritten := make(map[string]bool) + for _, d := range s.Data { + if err := d.writeTo(cw, s.When, options, metricsWritten); err != nil { + return cw.Written(), err + } + } + if _, err := io.WriteString(cw, fmt.Sprintf("# End of metric snapshot taken at: %v\n\n", s.When)); err != nil { + return cw.Written(), err + } + if err := cw.w.Flush(); err != nil { + return cw.Written(), err + } + return cw.Written(), nil +} diff --git a/pkg/prometheus/prometheus_test.go b/pkg/prometheus/prometheus_test.go new file mode 100644 index 000000000..95c6c3d18 --- /dev/null +++ b/pkg/prometheus/prometheus_test.go @@ -0,0 +1,653 @@ +// Copyright 2022 The gVisor Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package prometheus + +import ( + "bytes" + "errors" + "fmt" + "math" + "strings" + "sync" + "testing" + "time" + + v1proto "github.com/golang/protobuf/proto" + "github.com/google/go-cmp/cmp" + "github.com/prometheus/common/expfmt" + "google.golang.org/protobuf/encoding/prototext" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/testing/protocmp" + pb "gvisor.dev/gvisor/pkg/metric/metric_go_proto" +) + +// timeNowMu is used to synchronize injection of time.Now. +var timeNowMu sync.Mutex + +// at executes a function with the clock returning a given time. +func at(when time.Time, f func()) { + timeNowMu.Lock() + defer timeNowMu.Unlock() + previousFunc := timeNow + timeNow = func() time.Time { return when } + defer func() { timeNow = previousFunc }() + f() +} + +// newSnapshotAt creates a new Snapshot with the given timestamp. +func newSnapshotAt(when time.Time) *Snapshot { + var s *Snapshot + at(when, func() { + s = NewSnapshot() + }) + return s +} + +// Helper builder type for metric metadata. +type metricMetadata struct { + PB *pb.MetricMetadata + Fields map[string]string +} + +func (m *metricMetadata) clone() *metricMetadata { + m2 := &metricMetadata{ + PB: &pb.MetricMetadata{}, + Fields: make(map[string]string, len(m.Fields)), + } + proto.Merge(m2.PB, m.PB) + for k, v := range m.Fields { + m2.Fields[k] = v + } + return m2 +} + +// metric returns the Metric metadata struct for this metric metadata. +func (m *metricMetadata) metric() *Metric { + var metricType Type + switch m.PB.GetType() { + case pb.MetricMetadata_TYPE_UINT64: + if m.PB.GetCumulative() { + metricType = TypeCounter + } else { + metricType = TypeGauge + } + case pb.MetricMetadata_TYPE_DISTRIBUTION: + metricType = TypeHistogram + default: + panic(fmt.Sprintf("invalid type %v", m.PB.GetType())) + } + return &Metric{ + Name: m.PB.GetPrometheusName(), + Type: metricType, + Help: m.PB.GetDescription(), + } +} + +// Convenient metric field metadata definitions. +var ( + field1 = &pb.MetricMetadata_Field{ + FieldName: "field1", + AllowedValues: []string{"val1a", "val1b"}, + } + field2 = &pb.MetricMetadata_Field{ + FieldName: "field2", + AllowedValues: []string{"val2a", "val2b"}, + } +) + +// fieldVal returns a copy of this *metricMetadata with the given field-value +// stored on the side of the metadata. Meant to be used during snapshot data +// construction, where methods like int() make it easy to construct *Data +// structs with field values. +func (m *metricMetadata) fieldVal(field *pb.MetricMetadata_Field, val string) *metricMetadata { + return m.fieldVals(map[*pb.MetricMetadata_Field]string{field: val}) +} + +// fieldVals acts like fieldVal but for multiple fields, at the expense of +// having a less convenient function signature. +func (m *metricMetadata) fieldVals(fieldToVal map[*pb.MetricMetadata_Field]string) *metricMetadata { + m2 := m.clone() + for field, val := range fieldToVal { + m2.Fields[field.GetFieldName()] = val + } + return m2 +} + +// labels returns a label key-value map associated with the metricMetadata. +func (m *metricMetadata) labels() map[string]string { + if len(m.Fields) == 0 { + return nil + } + return m.Fields +} + +// int returns a new Data struct with the given value for the current metric. +// If the current metric has fields, all of its fields must accept exactly one +// value, and this value will be used as the value for that field. +// If a field accepts multiple values, the function will panic. +func (m *metricMetadata) int(val int64) *Data { + data := NewIntData(m.metric(), val) + data.Labels = m.labels() + return data +} + +// float returns a new Data struct with the given value for the current metric. +// If the current metric has fields, all of its fields must accept exactly one +// value, and this value will be used as the value for that field. +// If a field accepts multiple values, the function will panic. +func (m *metricMetadata) float(val float64) *Data { + data := NewFloatData(m.metric(), val) + data.Labels = m.labels() + return data +} + +// float returns a new Data struct with the given value for the current metric. +// If the current metric has fields, all of its fields must accept exactly one +// value, and this value will be used as the value for that field. +// If a field accepts multiple values, the function will panic. +func (m *metricMetadata) dist(samples ...int64) *Data { + var total int64 + buckets := make([]Bucket, len(m.PB.GetDistributionBucketLowerBounds())+1) + var bucket *Bucket + for i, lowerBound := range m.PB.GetDistributionBucketLowerBounds() { + (&buckets[i]).UpperBound = Number{Int: lowerBound} + } + (&buckets[len(buckets)-1]).UpperBound = Number{Float: math.Inf(1)} + for _, sample := range samples { + total += sample + bucket = &buckets[0] + for i, lowerBound := range m.PB.GetDistributionBucketLowerBounds() { + if sample >= lowerBound { + bucket = &buckets[i+1] + } else { + break + } + } + bucket.Samples++ + } + return &Data{ + Metric: m.metric(), + Labels: m.labels(), + HistogramValue: &Histogram{ + Total: Number{Int: total}, + Buckets: buckets, + }, + } +} + +// Convenient metric metadata definitions. +var ( + fooInt = &metricMetadata{ + PB: &pb.MetricMetadata{ + Name: "fooInt", + PrometheusName: "foo_int", + Description: "An integer about foo", + Cumulative: false, + Units: pb.MetricMetadata_UNITS_NONE, + Sync: true, + Type: pb.MetricMetadata_TYPE_UINT64, + }, + } + fooCounter = &metricMetadata{ + PB: &pb.MetricMetadata{ + Name: "fooCounter", + PrometheusName: "foo_counter", + Description: "A counter of foos", + Cumulative: true, + Units: pb.MetricMetadata_UNITS_NONE, + Sync: true, + Type: pb.MetricMetadata_TYPE_UINT64, + }, + } + fooDist = &metricMetadata{ + PB: &pb.MetricMetadata{ + Name: "fooDist", + PrometheusName: "foo_dist", + Description: "A distribution about foo", + Cumulative: false, + Units: pb.MetricMetadata_UNITS_NONE, + Sync: true, + Type: pb.MetricMetadata_TYPE_DISTRIBUTION, + DistributionBucketLowerBounds: []int64{0, 1, 2, 4, 8}, + }, + } +) + +// shortWriter implements io.Writer but fails after a given number of bytes. +type shortWriter struct { + buf bytes.Buffer + size int + maxSize int +} + +// Reset erases buffer data and resets the shortWriter to the given size. +func (s *shortWriter) Reset(size int) { + s.buf.Reset() + s.size = 0 + s.maxSize = size +} + +// String returns the buffered data as a string. +func (s *shortWriter) String() string { + return s.buf.String() +} + +// Write implements io.Writer.Write. +func (s *shortWriter) Write(b []byte) (n int, err error) { + toWrite := len(b) + leftToWrite := s.maxSize - s.size + if leftToWrite < toWrite { + toWrite = leftToWrite + } + if toWrite == 0 { + return 0, errors.New("writer out of capacity") + } + written, err := s.buf.Write(b[:toWrite]) + s.size += written + if written == len(b) { + return written, err + } + return written, errors.New("short write") +} + +// reflectProto converts a v1 or v2 proto message to a proto message with +// reflection enabled. +func reflectProto(m any) protoreflect.ProtoMessage { + if msg, hasReflection := m.(proto.Message); hasReflection { + return msg + } + // Convert v1 proto to introspectable view, if possible and necessary. + if v1pb, ok := m.(v1proto.Message); ok { + return v1proto.MessageReflect(v1pb).Interface() + } + panic(fmt.Sprintf("Proto message %v isn't of a supported protobuf type", m)) +} + +// TestSnapshotToPrometheus verifies that the contents of a Snapshot can be +// converted into text that can be parsed by the Prometheus parsing libraries, +// and produces the data we expect them to. +func TestSnapshotToPrometheus(t *testing.T) { + singleLineFormatter := &prototext.MarshalOptions{Multiline: false, EmitUnknown: true} + multiLineFormatter := &prototext.MarshalOptions{Multiline: true, Indent: " ", EmitUnknown: true} + testStart := time.Now() + newSnapshot := func() *Snapshot { + return newSnapshotAt(testStart) + } + for _, test := range []struct { + Name string + + // Snapshot will be rendered as Prometheus and compared against WantData. + Snapshot *Snapshot + + // ExportOptions dictates the options used during Snapshot rendering. + ExportOptions ExportOptions + + // WantFail, if true, indicates that the test is expected to fail when + // rendering or parsing the snapshot data. + WantFail bool + + // WantData is Prometheus text format that matches the data in Snapshot. + // The substring "{TIMESTAMP}" will be replaced with the value of + // `testStart` in milliseconds. + WantData string + }{ + { + Name: "empty snapshot", + Snapshot: newSnapshot(), + }, + { + Name: "simple integer", + Snapshot: newSnapshot().Add(fooInt.int(3)), + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int 3 {TIMESTAMP} + `, + }, + { + Name: "simple float", + Snapshot: newSnapshot().Add(fooInt.float(2.5)), + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int 2.5 {TIMESTAMP} + `, + }, + { + Name: "simple counter", + Snapshot: newSnapshot().Add(fooCounter.int(4)), + WantData: ` + # HELP foo_counter A counter of foos + # TYPE foo_counter counter + foo_counter 4 {TIMESTAMP} + `, + }, + { + Name: "two metrics", + Snapshot: newSnapshot().Add( + // Note the different order here than in WantData, + // to test ordering independence. + fooCounter.int(4), + fooInt.int(3), + ), + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int 3 {TIMESTAMP} + # HELP foo_counter A counter of foos + # TYPE foo_counter counter + foo_counter 4 {TIMESTAMP} + `, + }, + { + Name: "metric with 1 field", + Snapshot: newSnapshot().Add( + fooInt.fieldVal(field1, "val1a").int(3), + fooInt.fieldVal(field1, "val1b").int(7), + ), + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int{field1="val1a"} 3 {TIMESTAMP} + foo_int{field1="val1b"} 7 {TIMESTAMP} + `, + }, + { + Name: "metric with 2 fields", + Snapshot: newSnapshot().Add( + fooInt.fieldVal(field1, "val1a").fieldVal(field2, "val2a").int(3), + fooInt.fieldVal(field2, "val2b").fieldVal(field1, "val1b").int(7), + ), + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int{field1="val1a",field2="val2a"} 3 {TIMESTAMP} + foo_int{field1="val1b",field2="val2b"} 7 {TIMESTAMP} + `, + }, + { + Name: "simple integer with export options", + Snapshot: newSnapshot().Add(fooInt.int(3)), + ExportOptions: ExportOptions{ + CommentHeader: "Some header", + ExporterPrefix: "some_prefix_", + ExtraLabels: map[string]string{ + "field3": "val3a", + }, + }, + WantData: ` + # HELP some_prefix_foo_int An integer about foo + # TYPE some_prefix_foo_int gauge + some_prefix_foo_int{field3="val3a"} 3 {TIMESTAMP} + `, + }, + { + Name: "integer with fields mixing with export options", + Snapshot: newSnapshot().Add( + fooInt.fieldVal(field1, "val1a").fieldVal(field2, "val2a").int(3), + fooInt.fieldVal(field2, "val2b").fieldVal(field1, "val1b").int(7), + ), + ExportOptions: ExportOptions{ + ExtraLabels: map[string]string{ + "field3": "val3a", + }, + }, + WantData: ` + # HELP foo_int An integer about foo + # TYPE foo_int gauge + foo_int{field1="val1a",field2="val2a",field3="val3a"} 3 {TIMESTAMP} + foo_int{field1="val1b",field2="val2b",field3="val3a"} 7 {TIMESTAMP} + `, + }, + { + Name: "integer with fields conflicting with export options", + Snapshot: newSnapshot().Add( + fooInt.fieldVal(field1, "val1a").fieldVal(field2, "val2a").int(3), + fooInt.fieldVal(field2, "val2b").fieldVal(field1, "val1b").int(7), + ), + ExportOptions: ExportOptions{ + ExtraLabels: map[string]string{ + "field2": "val2c", + "field3": "val3a", + }, + }, + WantFail: true, + }, + { + Name: "simple distribution", + Snapshot: newSnapshot().Add( + // -1 + 3 + 3 + 3 + 5 + 7 + 7 + 99 = 126 + fooDist.dist(-1, 3, 3, 3, 5, 7, 7, 99), + ), + WantData: ` + # HELP foo_dist A distribution about foo + # TYPE foo_dist histogram + foo_dist_bucket{le="0"} 1 {TIMESTAMP} + foo_dist_bucket{le="1"} 1 {TIMESTAMP} + foo_dist_bucket{le="2"} 1 {TIMESTAMP} + foo_dist_bucket{le="4"} 4 {TIMESTAMP} + foo_dist_bucket{le="8"} 7 {TIMESTAMP} + foo_dist_bucket{le="+inf"} 8 {TIMESTAMP} + foo_dist_sum 126 {TIMESTAMP} + foo_dist_count 8 {TIMESTAMP} + `, + }, + { + Name: "distribution with 'le' label", + Snapshot: newSnapshot().Add( + fooDist.fieldVal(&pb.MetricMetadata_Field{ + FieldName: "le", + AllowedValues: []string{"foo"}, + }, "foo").dist(-1, 3, 3, 3, 5, 7, 7, 99), + ), + WantFail: true, + }, + { + Name: "distribution with no samples", + Snapshot: newSnapshot().Add( + fooDist.dist(), + ), + WantData: ` + # HELP foo_dist A distribution about foo + # TYPE foo_dist histogram + foo_dist_bucket{le="0"} 0 {TIMESTAMP} + foo_dist_bucket{le="1"} 0 {TIMESTAMP} + foo_dist_bucket{le="2"} 0 {TIMESTAMP} + foo_dist_bucket{le="4"} 0 {TIMESTAMP} + foo_dist_bucket{le="8"} 0 {TIMESTAMP} + foo_dist_bucket{le="+inf"} 0 {TIMESTAMP} + foo_dist_sum 0 {TIMESTAMP} + foo_dist_count 0 {TIMESTAMP} + `, + }, + { + Name: "distribution with 1 field", + Snapshot: newSnapshot().Add( + // -1 + 3 + 3 + 3 + 5 + 7 + 7 + 99 = 126 + fooDist.fieldVal(field1, "val1a").dist(-1, 3, 3, 3, 5, 7, 7, 99), + // 3 + 5 + 3 = 11 + fooDist.fieldVal(field1, "val1b").dist(3, 5, 3), + ), + WantData: ` + # HELP foo_dist A distribution about foo + # TYPE foo_dist histogram + foo_dist_bucket{field1="val1a",le="0"} 1 {TIMESTAMP} + foo_dist_bucket{field1="val1a",le="1"} 1 {TIMESTAMP} + foo_dist_bucket{field1="val1a",le="2"} 1 {TIMESTAMP} + foo_dist_bucket{field1="val1a",le="4"} 4 {TIMESTAMP} + foo_dist_bucket{field1="val1a",le="8"} 7 {TIMESTAMP} + foo_dist_bucket{field1="val1a",le="+inf"} 8 {TIMESTAMP} + foo_dist_sum{field1="val1a"} 126 {TIMESTAMP} + foo_dist_count{field1="val1a"} 8 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="0"} 0 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="1"} 0 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="2"} 0 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="4"} 2 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="8"} 3 {TIMESTAMP} + foo_dist_bucket{field1="val1b",le="+inf"} 3 {TIMESTAMP} + foo_dist_sum{field1="val1b"} 11 {TIMESTAMP} + foo_dist_count{field1="val1b"} 3 {TIMESTAMP} + `, + }, + { + Name: "distribution with 2 fields, one from ExportOptions", + Snapshot: newSnapshot().Add( + // -1 + 3 + 3 + 3 + 5 + 7 + 7 + 99 = 126 + fooDist.fieldVal(field1, "val1a").dist(-1, 3, 3, 3, 5, 7, 7, 99), + // 3 + 5 + 3 = 11 + fooDist.fieldVal(field1, "val1b").dist(3, 5, 3), + ), + ExportOptions: ExportOptions{ + CommentHeader: "Some header", + ExporterPrefix: "some_prefix_", + ExtraLabels: map[string]string{"field2": "val2a"}, + }, + WantData: ` + # HELP some_prefix_foo_dist A distribution about foo + # TYPE some_prefix_foo_dist histogram + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="0"} 1 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="1"} 1 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="2"} 1 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="4"} 4 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="8"} 7 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1a",field2="val2a",le="+inf"} 8 {TIMESTAMP} + some_prefix_foo_dist_sum{field1="val1a",field2="val2a"} 126 {TIMESTAMP} + some_prefix_foo_dist_count{field1="val1a",field2="val2a"} 8 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="0"} 0 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="1"} 0 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="2"} 0 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="4"} 2 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="8"} 3 {TIMESTAMP} + some_prefix_foo_dist_bucket{field1="val1b",field2="val2a",le="+inf"} 3 {TIMESTAMP} + some_prefix_foo_dist_sum{field1="val1b",field2="val2a"} 11 {TIMESTAMP} + some_prefix_foo_dist_count{field1="val1b",field2="val2a"} 3 {TIMESTAMP} + `, + }, + } { + t.Run(test.Name, func(t *testing.T) { + // Render and parse snapshot data. + var buf bytes.Buffer + if _, err := test.Snapshot.WriteTo(&buf, test.ExportOptions); err != nil { + if test.WantFail { + return + } + t.Fatalf("cannot write snapshot: %v", err) + } + gotMetricsRaw := buf.String() + gotMetrics, err := (&expfmt.TextParser{}).TextToMetricFamilies(&buf) + if err != nil { + if test.WantFail { + return + } + t.Fatalf("cannot parse data written from snapshot: %v", err) + } + if test.WantFail { + t.Fatalf("Test unexpectedly succeeded to render and parse snapshot data") + } + + // Verify that the data is consistent (i.e. verify that it's not based on random map ordering) + var buf2 bytes.Buffer + if _, err := test.Snapshot.WriteTo(&buf2, test.ExportOptions); err != nil { + if test.WantFail { + return + } + t.Fatalf("cannot write snapshot: %v", err) + } + gotMetricsRaw2 := buf2.String() + if gotMetricsRaw != gotMetricsRaw2 { + t.Errorf("inconsistent snapshot rendering:\n\n%s\n\n---- VS ----\n\n%s\n\n", gotMetricsRaw, gotMetricsRaw2) + } + + // Verify that error propagation works by having the writer fail at each possible spot. + // This exercises all the write error propagation branches. + var shortWriter shortWriter + for writeLength := 0; writeLength < len(gotMetricsRaw); writeLength++ { + shortWriter.Reset(writeLength) + if _, err := test.Snapshot.WriteTo(&shortWriter, test.ExportOptions); err == nil { + t.Fatalf("snapshot data unexpectedly succeeded being written to short writer (length %d): %v", writeLength, shortWriter.String()) + } + if shortWriter.size != writeLength { + t.Fatalf("Short writer should have allowed %d bytes of snapshot data to be written, but actual number of bytes written is %d bytes", writeLength, shortWriter.size) + } + } + + // Parse reference data. + wantData := strings.ReplaceAll(test.WantData, "{TIMESTAMP}", fmt.Sprintf("%d", testStart.UnixMilli())) + wantMetrics, err := (&expfmt.TextParser{}).TextToMetricFamilies(strings.NewReader(wantData)) + if err != nil { + t.Fatalf("cannot parse reference data: %v", err) + } + + if len(test.Snapshot.Data) != 0 { + // If the snapshot isn't empty, verify that the data we got from both `got` and `want` + // is non-zero. Otherwise, this whole test could accidentally succeed by having all attempts + // at parsing the data result into an empty set. + if len(wantMetrics) == 0 { + t.Error("Snapshot is not empty, but parsing the reference data resulted in no data being produced") + } + if len(gotMetrics) == 0 { + t.Error("Snapshot is not empty, but parsing the rendered snapshot resulted in no data being produced") + } + } + + // Verify that all of `wantMetrics` is in `gotMetrics`. + for metric, want := range wantMetrics { + if _, found := gotMetrics[metric]; !found { + wantText, err := singleLineFormatter.Marshal(reflectProto(want)) + if err != nil { + t.Fatalf("cannot marshal reference data: %v", err) + } + t.Errorf("metric %s is in reference data (%v) but not present in snapshot data", metric, string(wantText)) + } + } + + // Verify that all of `gotMetrics` is in `wantMetrics`. + for metric, got := range gotMetrics { + if _, found := wantMetrics[metric]; !found { + gotText, err := singleLineFormatter.Marshal(reflectProto(got)) + if err != nil { + t.Fatalf("cannot marshal snapshot data: %v", err) + } + t.Errorf("metric %s found in snapshot data (%v) but not present in reference data", metric, string(gotText)) + } + } + + // The rest of the test assumes the keys are the same. + if t.Failed() { + return + } + + // Verify metric data matches. + for metric := range wantMetrics { + t.Run(metric, func(t *testing.T) { + want := reflectProto(wantMetrics[metric]) + got := reflectProto(gotMetrics[metric]) + if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" { + wantText, err := multiLineFormatter.Marshal(want) + if err != nil { + t.Fatalf("cannot marshal reference data: %v", err) + } + gotText, err := multiLineFormatter.Marshal(got) + if err != nil { + t.Fatalf("cannot marshal snapshot data: %v", err) + } + t.Errorf("Snapshot data did not produce the same data as the reference data.\n\nReference data:\n\n%v\n\nSnapshot data:\n\n%v\n\nDiff:\n\n%v\n\n", string(wantText), string(gotText), diff) + } + }) + } + }) + } +}