gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/prometheus/prometheus.go (about)

     1  // Copyright 2022 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package prometheus contains Prometheus-compliant metric data structures and utilities.
    16  // It can export data in Prometheus data format, documented at:
    17  // https://prometheus.io/docs/instrumenting/exposition_formats/
    18  package prometheus
    19  
    20  import (
    21  	"bytes"
    22  	"errors"
    23  	"fmt"
    24  	"io"
    25  	"math"
    26  	"reflect"
    27  	"sort"
    28  	"strings"
    29  	"time"
    30  )
    31  
    32  // timeNow is the time.Now() function. Can be mocked in tests.
    33  var timeNow = time.Now
    34  
    35  // Prometheus label names used to identify each sandbox.
    36  const (
    37  	SandboxIDLabel   = "sandbox"
    38  	PodNameLabel     = "pod_name"
    39  	NamespaceLabel   = "namespace_name"
    40  	IterationIDLabel = "iteration"
    41  )
    42  
    43  // Type is a Prometheus metric type.
    44  type Type int
    45  
    46  // List of supported Prometheus metric types.
    47  const (
    48  	TypeUntyped = Type(iota)
    49  	TypeGauge
    50  	TypeCounter
    51  	TypeHistogram
    52  )
    53  
    54  // Metric is a Prometheus metric metadata.
    55  type Metric struct {
    56  	// Name is the Prometheus metric name.
    57  	Name string `json:"name"`
    58  
    59  	// Type is the type of the metric.
    60  	Type Type `json:"type"`
    61  
    62  	// Help is an optional helpful string explaining what the metric is about.
    63  	Help string `json:"help"`
    64  }
    65  
    66  // writeMetricHeaderTo writes the metric comment header to the given writer.
    67  func writeMetricHeaderTo[T io.StringWriter](w T, m *Metric, options SnapshotExportOptions) error {
    68  	if m.Help != "" {
    69  		// This writes each string component one by one (rather than using fmt.Sprintf)
    70  		// in order to avoid allocating strings for each metric.
    71  		if _, err := w.WriteString("# HELP "); err != nil {
    72  			return err
    73  		}
    74  		if _, err := w.WriteString(options.ExporterPrefix); err != nil {
    75  			return err
    76  		}
    77  		if _, err := w.WriteString(m.Name); err != nil {
    78  			return err
    79  		}
    80  		if _, err := w.WriteString(" "); err != nil {
    81  			return err
    82  		}
    83  		if _, err := writeEscapedString(w, m.Help, false); err != nil {
    84  			return err
    85  		}
    86  		if _, err := w.WriteString("\n"); err != nil {
    87  			return err
    88  		}
    89  	}
    90  	var metricType string
    91  	switch m.Type {
    92  	case TypeGauge:
    93  		metricType = "gauge"
    94  	case TypeCounter:
    95  		metricType = "counter"
    96  	case TypeHistogram:
    97  		metricType = "histogram"
    98  	case TypeUntyped:
    99  		metricType = "untyped"
   100  	}
   101  	if metricType != "" {
   102  		if _, err := w.WriteString("# TYPE "); err != nil {
   103  			return err
   104  		}
   105  		if _, err := w.WriteString(options.ExporterPrefix); err != nil {
   106  			return err
   107  		}
   108  		if _, err := w.WriteString(m.Name); err != nil {
   109  			return err
   110  		}
   111  		if _, err := w.WriteString(" "); err != nil {
   112  			return err
   113  		}
   114  		if _, err := w.WriteString(metricType); err != nil {
   115  			return err
   116  		}
   117  		if _, err := w.WriteString("\n"); err != nil {
   118  			return err
   119  		}
   120  	}
   121  	return nil
   122  }
   123  
   124  // Number represents a numerical value.
   125  // In Prometheus, all numbers are float64s.
   126  // However, for the purpose of usage of this library, we support expressing numbers as integers,
   127  // which makes things like counters much easier and more precise.
   128  // At data export time (i.e. when written out in Prometheus data format), it is coalesced into
   129  // a float.
   130  type Number struct {
   131  	// Float is the float value of this number.
   132  	// Mutually exclusive with Int.
   133  	Float float64 `json:"float,omitempty"`
   134  
   135  	// Int is the integer value of this number.
   136  	// Mutually exclusive with Float.
   137  	Int int64 `json:"int,omitempty"`
   138  }
   139  
   140  // Common numbers which are reused and don't need their own memory allocations.
   141  var (
   142  	zero        = Number{}
   143  	intOne      = Number{Int: 1}
   144  	floatOne    = Number{Float: 1.0}
   145  	floatNaN    = Number{Float: math.NaN()}
   146  	floatInf    = Number{Float: math.Inf(1)}
   147  	floatNegInf = Number{Float: math.Inf(-1)}
   148  )
   149  
   150  // NewInt returns a new integer Number.
   151  func NewInt(val int64) *Number {
   152  	switch val {
   153  	case 0:
   154  		return &zero
   155  	case 1:
   156  		return &intOne
   157  	default:
   158  		return &Number{Int: val}
   159  	}
   160  }
   161  
   162  // NewFloat returns a new floating-point Number.
   163  func NewFloat(val float64) *Number {
   164  	if math.IsNaN(val) {
   165  		return &floatNaN
   166  	}
   167  	switch val {
   168  	case 0:
   169  		return &zero
   170  	case 1.0:
   171  		return &floatOne
   172  	case math.Inf(1.0):
   173  		return &floatInf
   174  	case math.Inf(-1.0):
   175  		return &floatNegInf
   176  	default:
   177  		return &Number{Float: val}
   178  	}
   179  }
   180  
   181  // IsInteger returns whether this number contains an integer value.
   182  // This is defined as either having the `Float` part set to zero (in which case the `Int` part takes
   183  // precedence), or having `Float` be a value equal to its own rounding and not a special float.
   184  //
   185  //go:nosplit
   186  func (n *Number) IsInteger() bool {
   187  	if n.Float == 0 {
   188  		return true
   189  	}
   190  	if math.IsNaN(n.Float) || n.Float == math.Inf(-1) || n.Float == math.Inf(1) {
   191  		return false
   192  	}
   193  	return n.Float < float64(math.MaxInt64) && n.Float > float64(math.MinInt64) && math.Round(n.Float) == n.Float
   194  }
   195  
   196  // ToFloat returns this number as a floating-point number, regardless of which
   197  // type the number was encoded as. An integer Number will have its value cast
   198  // to a float, while a floating-point Number will have its value returned
   199  // as-is.
   200  //
   201  //go:nosplit
   202  func (n *Number) ToFloat() float64 {
   203  	if n.Int != 0 {
   204  		return float64(n.Int)
   205  	}
   206  	return n.Float
   207  }
   208  
   209  // String returns a string representation of this number.
   210  func (n *Number) String() string {
   211  	var s strings.Builder
   212  	if err := writeNumberTo(&s, n); err != nil {
   213  		panic(err)
   214  	}
   215  	return s.String()
   216  }
   217  
   218  // SameType returns true if `n` and `other` are either both floating-point or both integers.
   219  // If a `Number` is zero, it is considered of the same type as any other zero `Number`.
   220  //
   221  //go:nosplit
   222  func (n *Number) SameType(other *Number) bool {
   223  	// Within `n` and `other`, at least one of `Int` or `Float` must be set to zero.
   224  	// Therefore, this verifies that there is at least one shared zero between the two.
   225  	return n.Float == other.Float || n.Int == other.Int
   226  }
   227  
   228  // GreaterThan returns true if n > other.
   229  // Precondition: n.SameType(other) is true. Panics otherwise.
   230  //
   231  //go:nosplit
   232  func (n *Number) GreaterThan(other *Number) bool {
   233  	if !n.SameType(other) {
   234  		panic("tried to compare two numbers of different types")
   235  	}
   236  	if n.IsInteger() {
   237  		return n.Int > other.Int
   238  	}
   239  	return n.Float > other.Float
   240  }
   241  
   242  // WriteInteger writes the given integer to a writer without allocating strings.
   243  //
   244  //go:nosplit
   245  func WriteInteger[T io.StringWriter](w T, val int64) (int, error) {
   246  	const decimalDigits = "0123456789"
   247  	if val == 0 {
   248  		return w.WriteString(decimalDigits[0:1])
   249  	}
   250  	var written int
   251  	if val < 0 {
   252  		n, err := w.WriteString("-")
   253  		written += n
   254  		if err != nil {
   255  			return written, err
   256  		}
   257  		val = -val
   258  	}
   259  	decimal := int64(1)
   260  	for ; val/decimal != 0; decimal *= 10 {
   261  	}
   262  	for decimal /= 10; decimal > 0; decimal /= 10 {
   263  		digit := (val / decimal) % 10
   264  		n, err := w.WriteString(decimalDigits[digit : digit+1])
   265  		written += n
   266  		if err != nil {
   267  			return written, err
   268  		}
   269  	}
   270  	return written, nil
   271  }
   272  
   273  // WriteHex writes the given integer as hex to a writer
   274  // without allocating strings.
   275  //
   276  //go:nosplit
   277  func WriteHex[T io.StringWriter](w T, val uint64) (int, error) {
   278  	const hexDigits = "0123456789abcdef"
   279  	if val == 0 {
   280  		return w.WriteString(hexDigits[0:1])
   281  	}
   282  	var written int
   283  	hex := uint64(16)
   284  	for ; val/hex != 0; hex <<= 4 {
   285  	}
   286  	for hex >>= 4; hex > 0; hex >>= 4 {
   287  		digit := (val / hex) % 16
   288  		n, err := w.WriteString(hexDigits[digit : digit+1])
   289  		written += n
   290  		if err != nil {
   291  			return written, err
   292  		}
   293  	}
   294  	return written, nil
   295  }
   296  
   297  // writeNumberTo writes the number to the given writer.
   298  // This only causes heap allocations when the number is a non-zero, non-special float.
   299  func writeNumberTo[T io.StringWriter](w T, n *Number) error {
   300  	var s string
   301  	switch {
   302  	// Zero case:
   303  	case n.Int == 0 && n.Float == 0:
   304  		s = "0"
   305  
   306  	// Integer case:
   307  	case n.Int != 0:
   308  		_, err := WriteInteger(w, n.Int)
   309  		return err
   310  
   311  	// Special float cases:
   312  	case n.Float == math.Inf(-1):
   313  		s = "-Inf"
   314  	case n.Float == math.Inf(1):
   315  		s = "+Inf"
   316  	case math.IsNaN(n.Float):
   317  		s = "NaN"
   318  
   319  	// Regular float case:
   320  	default:
   321  		s = fmt.Sprintf("%f", n.Float)
   322  	}
   323  	_, err := w.WriteString(s)
   324  	return err
   325  }
   326  
   327  // Bucket is a single histogram bucket.
   328  type Bucket struct {
   329  	// UpperBound is the upper bound of the bucket.
   330  	// The lower bound of the bucket is the largest UpperBound within other Histogram Buckets that
   331  	// is smaller than this bucket's UpperBound.
   332  	// The bucket with the smallest UpperBound within a Histogram implicitly has -Inf as lower bound.
   333  	// This should be set to +Inf to mark the "last" bucket.
   334  	UpperBound Number `json:"le"`
   335  
   336  	// Samples is the number of samples in the bucket.
   337  	// Note: When exported to Prometheus, they are exported cumulatively, i.e. the count of samples
   338  	// exported in Bucket i is actually sum(histogram.Buckets[j].Samples for 0 <= j <= i).
   339  	Samples uint64 `json:"n,omitempty"`
   340  }
   341  
   342  // Histogram contains data about histogram values.
   343  type Histogram struct {
   344  	// Total is the sum of sample values across all buckets.
   345  	Total Number `json:"total"`
   346  	// Min is the minimum sample ever recorded in this histogram.
   347  	Min Number `json:"min"`
   348  	// Max is the maximum sample ever recorded in this histogram.
   349  	Max Number `json:"max"`
   350  	// SumOfSquaredDeviations is the number of squared deviations of all samples.
   351  	SumOfSquaredDeviations Number `json:"ssd"`
   352  	// Buckets contains per-bucket data.
   353  	// A distribution with n finite-boundary buckets should have n+2 entries here.
   354  	// The 0th entry is the underflow bucket (i.e. the one with -inf as lower bound),
   355  	// and the last aka (n+1)th entry is the overflow bucket (i.e. the one with +inf as upper bound).
   356  	Buckets []Bucket `json:"buckets,omitempty"`
   357  }
   358  
   359  // Data is an observation of the value of a single metric at a certain point in time.
   360  type Data struct {
   361  	// Metric is the metric for which the value is being reported.
   362  	Metric *Metric `json:"metric"`
   363  
   364  	// Labels is a key-value pair representing the labels set on this metric.
   365  	// This may be merged with other labels during export.
   366  	Labels map[string]string `json:"labels,omitempty"`
   367  
   368  	// ExternalLabels are more labels merged together with `Labels`.
   369  	// They can be set using SetExternalLabels.
   370  	// They are useful in the case where a single Data needs labels from two sources:
   371  	// labels specific to this data point (which should be in `Labels`), and labels
   372  	// that are shared between multiple data points (stored in `ExternalLabels`).
   373  	// This avoids allocating unique `Labels` maps for each Data struct, when
   374  	// most of the actual labels would be shared between them.
   375  	ExternalLabels map[string]string `json:"external_labels,omitempty"`
   376  
   377  	// At most one of the fields below may be set.
   378  	// Which one depends on the type of the metric.
   379  
   380  	// Number is used for all numerical types.
   381  	Number *Number `json:"val,omitempty"`
   382  
   383  	// Histogram is used for histogram-typed metrics.
   384  	HistogramValue *Histogram `json:"histogram,omitempty"`
   385  }
   386  
   387  // NewIntData returns a new Data struct with the given metric and value.
   388  func NewIntData(metric *Metric, val int64) *Data {
   389  	return LabeledIntData(metric, nil, val)
   390  }
   391  
   392  // LabeledIntData returns a new Data struct with the given metric, labels, and value.
   393  func LabeledIntData(metric *Metric, labels map[string]string, val int64) *Data {
   394  	return &Data{Metric: metric, Labels: labels, Number: NewInt(val)}
   395  }
   396  
   397  // NewFloatData returns a new Data struct with the given metric and value.
   398  func NewFloatData(metric *Metric, val float64) *Data {
   399  	return LabeledFloatData(metric, nil, val)
   400  }
   401  
   402  // LabeledFloatData returns a new Data struct with the given metric, labels, and value.
   403  func LabeledFloatData(metric *Metric, labels map[string]string, val float64) *Data {
   404  	return &Data{Metric: metric, Labels: labels, Number: NewFloat(val)}
   405  }
   406  
   407  // SetExternalLabels sets d.ExternalLabels. See its docstring for more information.
   408  // Returns `d` for chainability.
   409  func (d *Data) SetExternalLabels(externalLabels map[string]string) *Data {
   410  	d.ExternalLabels = externalLabels
   411  	return d
   412  }
   413  
   414  // ExportOptions contains options that control how metric data is exported in Prometheus format.
   415  type ExportOptions struct {
   416  	// CommentHeader is prepended as a comment before any metric data is exported.
   417  	CommentHeader string
   418  
   419  	// MetricsWritten memoizes written metric preambles (help/type comments)
   420  	// by metric name.
   421  	// If specified, this map can be used to avoid duplicate preambles across multiple snapshots.
   422  	// Note that this map is modified in-place during the writing process.
   423  	MetricsWritten map[string]bool
   424  }
   425  
   426  // SnapshotExportOptions contains options that control how metric data is exported for an
   427  // individual Snapshot.
   428  type SnapshotExportOptions struct {
   429  	// ExporterPrefix is prepended to all metric names.
   430  	ExporterPrefix string
   431  
   432  	// ExtraLabels is added as labels for all metric values.
   433  	ExtraLabels map[string]string
   434  }
   435  
   436  // writeEscapedString writes the given string in quotation marks and with some characters escaped,
   437  // per Prometheus spec. It does this without string allocations.
   438  // If `quoted` is true, quote characters will surround the string, and quote characters within `s`
   439  // will also be escaped.
   440  func writeEscapedString[T io.StringWriter](w T, s string, quoted bool) (int, error) {
   441  	const (
   442  		quote            = '"'
   443  		backslash        = '\\'
   444  		newline          = '\n'
   445  		quoteStr         = `"`
   446  		escapedQuote     = `\\"`
   447  		escapedBackslash = "\\\\"
   448  		escapedNewline   = "\\\n"
   449  	)
   450  	written := 0
   451  	var n int
   452  	var err error
   453  	if quoted {
   454  		n, err = w.WriteString(quoteStr)
   455  		written += n
   456  		if err != nil {
   457  			return written, err
   458  		}
   459  	}
   460  	for _, r := range s {
   461  		switch r {
   462  		case quote:
   463  			if quoted {
   464  				n, err = w.WriteString(escapedQuote)
   465  			} else {
   466  				n, err = w.WriteString(quoteStr)
   467  			}
   468  		case backslash:
   469  			n, err = w.WriteString(escapedBackslash)
   470  		case newline:
   471  			n, err = w.WriteString(escapedNewline)
   472  		default:
   473  			n, err = w.WriteString(string(r))
   474  		}
   475  		written += n
   476  		if err != nil {
   477  			return written, err
   478  		}
   479  	}
   480  	if quoted {
   481  		n, err = w.WriteString(quoteStr)
   482  		written += n
   483  		if err != nil {
   484  			return written, err
   485  		}
   486  	}
   487  	return written, nil
   488  }
   489  
   490  // writeMetricPreambleTo writes the metric name to the writer. It may also
   491  // write unwritten help and type comments of the metric if they haven't been
   492  // written to the writer yet.
   493  func writeMetricPreambleTo[T io.StringWriter](w T, d *Data, options SnapshotExportOptions, metricsWritten map[string]bool) error {
   494  	// Metric header, if we haven't printed it yet.
   495  	if !metricsWritten[d.Metric.Name] {
   496  		// Extra newline before each preamble for aesthetic reasons.
   497  		if _, err := w.WriteString("\n"); err != nil {
   498  			return err
   499  		}
   500  		if err := writeMetricHeaderTo(w, d.Metric, options); err != nil {
   501  			return err
   502  		}
   503  		metricsWritten[d.Metric.Name] = true
   504  	}
   505  
   506  	// Metric name.
   507  	if options.ExporterPrefix != "" {
   508  		if _, err := w.WriteString(options.ExporterPrefix); err != nil {
   509  			return err
   510  		}
   511  	}
   512  	if _, err := w.WriteString(d.Metric.Name); err != nil {
   513  		return err
   514  	}
   515  	return nil
   516  }
   517  
   518  // keyVal is a key-value pair used in the function below.
   519  type keyVal struct{ Key, Value string }
   520  
   521  // sortedIterateLabels iterates through labels and outputs them to `out` in sorted key order,
   522  // or stops when cancelCh is written to. It runs in O(n^2) time but makes no heap allocations.
   523  func sortedIterateLabels(labels map[string]string, out chan<- keyVal, cancelCh <-chan struct{}) {
   524  	defer close(out)
   525  	if len(labels) == 0 {
   526  		return
   527  	}
   528  
   529  	// smallestKey is the smallest key that we've already sent to `out`.
   530  	// It starts as the empty string, which means we haven't sent anything to `out` yet.
   531  	smallestKey := ""
   532  	// Find the smallest key of the whole set and send it out.
   533  	for k := range labels {
   534  		if smallestKey == "" || k < smallestKey {
   535  			smallestKey = k
   536  		}
   537  	}
   538  	select {
   539  	case out <- keyVal{smallestKey, labels[smallestKey]}:
   540  	case <-cancelCh:
   541  		return
   542  	}
   543  
   544  	// Iterate until we've sent as many items as we have as input to the output channel.
   545  	// We start at 1 because the loop above already sent out the smallest key to `out`.
   546  	for numOutput := 1; numOutput < len(labels); numOutput++ {
   547  		// nextSmallestKey is the smallest key that is strictly larger than `smallestKey`.
   548  		nextSmallestKey := ""
   549  		for k := range labels {
   550  			if k > smallestKey && (nextSmallestKey == "" || k < nextSmallestKey) {
   551  				nextSmallestKey = k
   552  			}
   553  		}
   554  
   555  		// Update smallestKey and send it out.
   556  		smallestKey = nextSmallestKey
   557  		select {
   558  		case out <- keyVal{smallestKey, labels[smallestKey]}:
   559  		case <-cancelCh:
   560  			return
   561  		}
   562  	}
   563  }
   564  
   565  // LabelOrError is used in OrderedLabels.
   566  // It represents either a key-value pair, or an error.
   567  type LabelOrError struct {
   568  	Key, Value string
   569  	Error      error
   570  }
   571  
   572  // OrderedLabels streams the list of 'label_key="label_value"' in sorted order, except "le" which is
   573  // a reserved Prometheus label name and should go last.
   574  // If an error is encountered, it is returned as the Error field of LabelOrError, and no further
   575  // messages will be sent on the channel.
   576  func OrderedLabels(labels ...map[string]string) <-chan LabelOrError {
   577  	// This function is quite hot on the metric-rendering path, and its naive "just put all the
   578  	// strings in one map to ensure no dupes it, then in one slice and sort it" approach is very
   579  	// allocation-heavy. This approach is more computation-heavy (it runs in
   580  	// O(len(labels) * len(largest label map))), but the only heap allocations it does is for the
   581  	// following tiny slices and channels. In practice, the number of label maps and the size of
   582  	// each label map is tiny, so this is worth doing despite the theoretically-longer run time.
   583  
   584  	// Initialize the channels we'll use.
   585  	mapChannels := make([]chan keyVal, 0, len(labels))
   586  	lastKeyVal := make([]keyVal, len(labels))
   587  	resultCh := make(chan LabelOrError)
   588  	var cancelCh chan struct{}
   589  	// outputError is a helper function for when we have encountered an error mid-way.
   590  	outputError := func(err error) {
   591  		if cancelCh != nil {
   592  			for range mapChannels {
   593  				cancelCh <- struct{}{}
   594  			}
   595  			close(cancelCh)
   596  		}
   597  		resultCh <- LabelOrError{Error: err}
   598  		close(resultCh)
   599  	}
   600  
   601  	// Verify that no label is the empty string. It's not a valid label name,
   602  	// and we use the empty string later on in the function as a marker of having
   603  	// finished processing all labels from a given label map.
   604  	for _, labelMap := range labels {
   605  		for label := range labelMap {
   606  			if label == "" {
   607  				go outputError(errors.New("got empty-string label"))
   608  				return resultCh
   609  			}
   610  		}
   611  	}
   612  
   613  	// Each label map is processed in its own goroutine,
   614  	// which will stream it back to this function in sorted order.
   615  	cancelCh = make(chan struct{}, len(labels))
   616  	for _, labelMap := range labels {
   617  		ch := make(chan keyVal)
   618  		mapChannels = append(mapChannels, ch)
   619  		go sortedIterateLabels(labelMap, ch, cancelCh)
   620  	}
   621  
   622  	// This goroutine is the meat of this function; it iterates through
   623  	// the results being streamed from each `sortedIterateLabels` goroutine
   624  	// that we spawned earlier, until all of them are exhausted or until we
   625  	// hit an error.
   626  	go func() {
   627  		// The "le" label is special and goes last, not in sorted order.
   628  		// gotLe is the empty string if there is no "le" label,
   629  		// otherwise it's the value of the "le" label.
   630  		var gotLe string
   631  
   632  		// numChannelsLeft tracks the number of channels that are still live.
   633  		for numChannelsLeft := len(mapChannels); numChannelsLeft > 0; {
   634  			// Iterate over all channels and ensure we have the freshest (smallest)
   635  			// label from each of them.
   636  			for i, ch := range mapChannels {
   637  				// A nil channel is one that has been closed.
   638  				if ch == nil {
   639  					continue
   640  				}
   641  				// If we already have the latest value from this channel,
   642  				// keep it there instead of getting a new one,
   643  				if lastKeyVal[i].Key != "" {
   644  					continue
   645  				}
   646  				// Otherwise, get a new label.
   647  				kv, open := <-ch
   648  				if !open {
   649  					// Channel has been closed, no more to read from this one.
   650  					numChannelsLeft--
   651  					mapChannels[i] = nil
   652  					continue
   653  				}
   654  				if kv.Key == "le" {
   655  					if gotLe != "" {
   656  						outputError(errors.New("got duplicate 'le' label"))
   657  						return
   658  					}
   659  					gotLe = kv.Value
   660  					continue
   661  				}
   662  				lastKeyVal[i] = kv
   663  			}
   664  
   665  			// We have one key-value pair from each still-active channel now.
   666  			// Find the smallest one between them.
   667  			smallestKey := ""
   668  			indexForSmallest := -1
   669  			for i, kv := range lastKeyVal {
   670  				if kv.Key == "" {
   671  					continue
   672  				}
   673  				if smallestKey == "" || kv.Key < smallestKey {
   674  					smallestKey = kv.Key
   675  					indexForSmallest = i
   676  				} else if kv.Key == smallestKey {
   677  					outputError(fmt.Errorf("got duplicate label %q", smallestKey))
   678  					return
   679  				}
   680  			}
   681  
   682  			if indexForSmallest == -1 {
   683  				// There are no more key-value pairs to output. We're done.
   684  				break
   685  			}
   686  
   687  			// Output the smallest key-value pairs out of all the channels.
   688  			resultCh <- LabelOrError{
   689  				Key:   smallestKey,
   690  				Value: lastKeyVal[indexForSmallest].Value,
   691  			}
   692  			// Mark the last key-value pair from the channel that gave us the
   693  			// smallest key-value pair as no longer present, so that we get a new
   694  			// key-value pair from it in the next iteration.
   695  			lastKeyVal[indexForSmallest] = keyVal{}
   696  		}
   697  
   698  		// Output the "le" label last.
   699  		if gotLe != "" {
   700  			resultCh <- LabelOrError{
   701  				Key:   "le",
   702  				Value: gotLe,
   703  			}
   704  		}
   705  		close(resultCh)
   706  		close(cancelCh)
   707  	}()
   708  
   709  	return resultCh
   710  }
   711  
   712  // writeLabelsTo writes a set of metric labels.
   713  func writeLabelsTo[T io.StringWriter](w T, d *Data, extraLabels map[string]string, leLabel *Number) error {
   714  	if len(d.Labels)+len(d.ExternalLabels)+len(extraLabels) != 0 || leLabel != nil {
   715  		if _, err := w.WriteString("{"); err != nil {
   716  			return err
   717  		}
   718  		var orderedLabels <-chan LabelOrError
   719  		if leLabel != nil {
   720  			orderedLabels = OrderedLabels(d.Labels, d.ExternalLabels, extraLabels, map[string]string{"le": leLabel.String()})
   721  		} else {
   722  			orderedLabels = OrderedLabels(d.Labels, d.ExternalLabels, extraLabels)
   723  		}
   724  		firstLabel := true
   725  		var foundError error
   726  		for labelOrError := range orderedLabels {
   727  			if foundError != nil {
   728  				continue
   729  			}
   730  			if labelOrError.Error != nil {
   731  				foundError = labelOrError.Error
   732  				continue
   733  			}
   734  			if !firstLabel {
   735  				if _, err := w.WriteString(","); err != nil {
   736  					return err
   737  				}
   738  			}
   739  			firstLabel = false
   740  			if _, err := w.WriteString(labelOrError.Key); err != nil {
   741  				return err
   742  			}
   743  			if _, err := w.WriteString("="); err != nil {
   744  				return err
   745  			}
   746  			if _, err := writeEscapedString(w, labelOrError.Value, true); err != nil {
   747  				return err
   748  			}
   749  		}
   750  		if foundError != nil {
   751  			return foundError
   752  		}
   753  		if _, err := w.WriteString("}"); err != nil {
   754  			return err
   755  		}
   756  	}
   757  	return nil
   758  }
   759  
   760  // writeMetricLine writes a single Data line with a single number (val) to w.
   761  func writeMetricLine[T io.StringWriter](w T, d *Data, metricSuffix string, val *Number, when time.Time, options SnapshotExportOptions, leLabel *Number, metricsWritten map[string]bool) error {
   762  	if err := writeMetricPreambleTo(w, d, options, metricsWritten); err != nil {
   763  		return err
   764  	}
   765  	if metricSuffix != "" {
   766  		if _, err := w.WriteString(metricSuffix); err != nil {
   767  			return err
   768  		}
   769  	}
   770  	if err := writeLabelsTo(w, d, options.ExtraLabels, leLabel); err != nil {
   771  		return err
   772  	}
   773  	if _, err := w.WriteString(" "); err != nil {
   774  		return err
   775  	}
   776  	if err := writeNumberTo(w, val); err != nil {
   777  		return err
   778  	}
   779  	if _, err := w.WriteString(" "); err != nil {
   780  		return err
   781  	}
   782  	if _, err := WriteInteger(w, when.UnixMilli()); err != nil {
   783  		return err
   784  	}
   785  	if _, err := w.WriteString("\n"); err != nil {
   786  		return err
   787  	}
   788  	return nil
   789  }
   790  
   791  // writeDataTo writes the Data to the given writer in Prometheus format.
   792  func writeDataTo[T io.StringWriter](w T, d *Data, when time.Time, options SnapshotExportOptions, metricsWritten map[string]bool) error {
   793  	switch d.Metric.Type {
   794  	case TypeUntyped, TypeGauge, TypeCounter:
   795  		return writeMetricLine(w, d, "", d.Number, when, options, nil, metricsWritten)
   796  	case TypeHistogram:
   797  		// Write an empty line before and after histograms to easily distinguish them from
   798  		// other metric lines.
   799  		if _, err := w.WriteString("\n"); err != nil {
   800  			return err
   801  		}
   802  		var numSamples uint64
   803  		var samples Number
   804  		for _, bucket := range d.HistogramValue.Buckets {
   805  			numSamples += bucket.Samples
   806  			samples.Int = int64(numSamples) // Prometheus distribution bucket counts are cumulative.
   807  			if err := writeMetricLine(w, d, "_bucket", &samples, when, options, &bucket.UpperBound, metricsWritten); err != nil {
   808  				return err
   809  			}
   810  		}
   811  		if err := writeMetricLine(w, d, "_sum", &d.HistogramValue.Total, when, options, nil, metricsWritten); err != nil {
   812  			return err
   813  		}
   814  		samples.Int = int64(numSamples)
   815  		if err := writeMetricLine(w, d, "_count", &samples, when, options, nil, metricsWritten); err != nil {
   816  			return err
   817  		}
   818  		if err := writeMetricLine(w, d, "_min", &d.HistogramValue.Min, when, options, nil, metricsWritten); err != nil {
   819  			return err
   820  		}
   821  		if err := writeMetricLine(w, d, "_max", &d.HistogramValue.Max, when, options, nil, metricsWritten); err != nil {
   822  			return err
   823  		}
   824  		if err := writeMetricLine(w, d, "_ssd", &d.HistogramValue.SumOfSquaredDeviations, when, options, nil, metricsWritten); err != nil {
   825  			return err
   826  		}
   827  		// Empty line after the histogram.
   828  		if _, err := w.WriteString("\n"); err != nil {
   829  			return err
   830  		}
   831  		return nil
   832  	default:
   833  		return fmt.Errorf("unknown metric type for metric %s: %v", d.Metric.Name, d.Metric.Type)
   834  	}
   835  }
   836  
   837  // Snapshot is a snapshot of the values of all the metrics at a certain point in time.
   838  type Snapshot struct {
   839  	// When is the timestamp at which the snapshot was taken.
   840  	// Note that Prometheus ultimately encodes timestamps as millisecond-precision int64s from epoch.
   841  	When time.Time `json:"when,omitempty"`
   842  
   843  	// Data is the whole snapshot data.
   844  	// Each Data must be a unique combination of (Metric, Labels) within a Snapshot.
   845  	Data []*Data `json:"data,omitempty"`
   846  }
   847  
   848  // NewSnapshot returns a new Snapshot at the current time.
   849  func NewSnapshot() *Snapshot {
   850  	return &Snapshot{When: timeNow()}
   851  }
   852  
   853  // Add data point(s) to the snapshot.
   854  // Returns itself for chainability.
   855  func (s *Snapshot) Add(data ...*Data) *Snapshot {
   856  	s.Data = append(s.Data, data...)
   857  	return s
   858  }
   859  
   860  const counterWriterBufSize = 32768
   861  
   862  // countingWriter implements io.StringWriter, and counts the number of bytes
   863  // written to it.
   864  // Useful in this file to keep track of total number of bytes without having
   865  // to plumb this everywhere in the writeX() functions in this file.
   866  type countingWriter[T io.StringWriter] struct {
   867  	buf        *bytes.Buffer
   868  	underlying T
   869  	written    int
   870  }
   871  
   872  // WriteString implements io.StringWriter.WriteString.
   873  // This avoids going into the slow, allocation-heavy path of io.WriteString.
   874  func (w *countingWriter[T]) WriteString(s string) (int, error) {
   875  	written, err := w.buf.WriteString(s)
   876  	w.written += written
   877  	if w.buf.Len() >= counterWriterBufSize {
   878  		w.Flush()
   879  	}
   880  	return written, err
   881  }
   882  
   883  func (w *countingWriter[T]) Flush() error {
   884  	if w.buf.Len() > 0 {
   885  		_, err := w.underlying.WriteString(w.buf.String())
   886  		w.buf.Reset()
   887  		return err
   888  	}
   889  	return nil
   890  }
   891  
   892  // Written returns the number of bytes written to the underlying writer (minus buffered writes).
   893  func (w *countingWriter[T]) Written() int {
   894  	return w.written - w.buf.Len()
   895  }
   896  
   897  // writeSnapshotSingleMetric writes a single metric data from a snapshot to
   898  // the given writer in Prometheus format.
   899  // It returns the number of bytes written.
   900  func writeSnapshotSingleMetric[T io.StringWriter](w T, s *Snapshot, options SnapshotExportOptions, metricName string, metricsWritten map[string]bool) error {
   901  	if !strings.HasPrefix(metricName, options.ExporterPrefix) {
   902  		return nil
   903  	}
   904  	wantMetricName := strings.TrimPrefix(metricName, options.ExporterPrefix)
   905  	for _, d := range s.Data {
   906  		if d.Metric.Name != wantMetricName {
   907  			continue
   908  		}
   909  		if err := writeDataTo(w, d, s.When, options, metricsWritten); err != nil {
   910  			return err
   911  		}
   912  	}
   913  	return nil
   914  }
   915  
   916  // ReusableWriter is a writer that can be reused to efficiently write
   917  // successive snapshots.
   918  type ReusableWriter[T io.StringWriter] struct {
   919  	// buf is the reusable buffer used for buffering writes.
   920  	// It is reset after each write, but keeps the underlying byte buffer,
   921  	// avoiding allocations on successive snapshot writes.
   922  	buf bytes.Buffer
   923  }
   924  
   925  // Write writes one or more snapshots to the writer.
   926  // This method may not be used concurrently for the same `ReusableWriter`.
   927  func (rw *ReusableWriter[T]) Write(w T, options ExportOptions, snapshotsToOptions map[*Snapshot]SnapshotExportOptions) (int, error) {
   928  	rw.buf.Reset()
   929  	cw := &countingWriter[T]{
   930  		buf:        &rw.buf,
   931  		underlying: w,
   932  	}
   933  	return write(cw, options, snapshotsToOptions)
   934  }
   935  
   936  // Write writes one or more snapshots to the writer.
   937  // This ensures same-name metrics across different snapshots are printed together, per spec.
   938  // If the caller will call `Write` successively for multiple snapshots, it is more efficient
   939  // to use the `ReusableWriter` type instead of this function.
   940  func Write[T io.StringWriter](w T, options ExportOptions, snapshotsToOptions map[*Snapshot]SnapshotExportOptions) (int, error) {
   941  	var b bytes.Buffer
   942  	// Sane default buffer size.
   943  	b.Grow(counterWriterBufSize)
   944  	cw := &countingWriter[T]{
   945  		buf:        &b,
   946  		underlying: w,
   947  	}
   948  	return write(cw, options, snapshotsToOptions)
   949  }
   950  
   951  func write[T io.StringWriter](cw *countingWriter[T], options ExportOptions, snapshotsToOptions map[*Snapshot]SnapshotExportOptions) (int, error) {
   952  	if len(snapshotsToOptions) == 0 {
   953  		return 0, nil
   954  	}
   955  	if options.CommentHeader != "" {
   956  		for _, commentLine := range strings.Split(options.CommentHeader, "\n") {
   957  			if _, err := cw.WriteString("# "); err != nil {
   958  				return cw.Written(), err
   959  			}
   960  			if _, err := cw.WriteString(commentLine); err != nil {
   961  				return cw.Written(), err
   962  			}
   963  			if _, err := cw.WriteString("\n"); err != nil {
   964  				return cw.Written(), err
   965  			}
   966  		}
   967  	}
   968  	snapshots := make([]*Snapshot, 0, len(snapshotsToOptions))
   969  	for snapshot := range snapshotsToOptions {
   970  		snapshots = append(snapshots, snapshot)
   971  	}
   972  	switch len(snapshots) {
   973  	case 1: // Single-snapshot case.
   974  		if _, err := cw.WriteString(fmt.Sprintf("# Writing data from snapshot containing %d data points taken at %v.\n", len(snapshots[0].Data), snapshots[0].When)); err != nil {
   975  			return cw.Written(), err
   976  		}
   977  	default: // Multi-snapshot case.
   978  		// Provide a consistent ordering of snapshots.
   979  		sort.Slice(snapshots, func(i, j int) bool {
   980  			return reflect.ValueOf(snapshots[i]).Pointer() < reflect.ValueOf(snapshots[j]).Pointer()
   981  		})
   982  		if _, err := cw.WriteString(fmt.Sprintf("# Writing data from %d snapshots:\n", len(snapshots))); err != nil {
   983  			return cw.Written(), err
   984  		}
   985  		for _, snapshot := range snapshots {
   986  			if _, err := cw.WriteString(fmt.Sprintf("#   - Snapshot with %d data points taken at %v: %v\n", len(snapshot.Data), snapshot.When, snapshotsToOptions[snapshot].ExtraLabels)); err != nil {
   987  				return cw.Written(), err
   988  			}
   989  		}
   990  	}
   991  	if _, err := cw.WriteString("\n"); err != nil {
   992  		return cw.Written(), err
   993  	}
   994  	if options.MetricsWritten == nil {
   995  		options.MetricsWritten = make(map[string]bool)
   996  	}
   997  	metricNamesMap := make(map[string]bool, len(options.MetricsWritten))
   998  	metricNames := make([]string, 0, len(options.MetricsWritten))
   999  	for _, snapshot := range snapshots {
  1000  		for _, data := range snapshot.Data {
  1001  			metricName := snapshotsToOptions[snapshot].ExporterPrefix + data.Metric.Name
  1002  			if !metricNamesMap[metricName] {
  1003  				metricNamesMap[metricName] = true
  1004  				metricNames = append(metricNames, metricName)
  1005  			}
  1006  		}
  1007  	}
  1008  	sort.Strings(metricNames)
  1009  	for _, metricName := range metricNames {
  1010  		for _, snapshot := range snapshots {
  1011  			writeSnapshotSingleMetric(cw, snapshot, snapshotsToOptions[snapshot], metricName, options.MetricsWritten)
  1012  		}
  1013  	}
  1014  	if _, err := cw.WriteString("\n# End of metric data.\n"); err != nil {
  1015  		return cw.Written(), err
  1016  	}
  1017  	if err := cw.Flush(); err != nil {
  1018  		return cw.Written(), err
  1019  	}
  1020  	return cw.Written(), nil
  1021  }