github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/metric/prometheus_exporter.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package metric
    12  
    13  import (
    14  	"io"
    15  
    16  	"github.com/gogo/protobuf/proto"
    17  	"github.com/prometheus/client_golang/prometheus"
    18  	prometheusgo "github.com/prometheus/client_model/go"
    19  	"github.com/prometheus/common/expfmt"
    20  )
    21  
    22  // PrometheusExporter contains a map of metric families (a metric with multiple labels).
    23  // It initializes each metric family once and reuses it for each prometheus scrape.
    24  // It is NOT thread-safe.
    25  // TODO(marc): we should really keep our metric objects here so we can avoid creating
    26  // new prometheus.Metric every time we are scraped.
    27  // see: https://github.com/cockroachdb/cockroach/issues/9326
    28  //  pe := MakePrometheusExporter()
    29  //  pe.AddMetricsFromRegistry(nodeRegistry)
    30  //  pe.AddMetricsFromRegistry(storeOneRegistry)
    31  //  ...
    32  //  pe.AddMetricsFromRegistry(storeNRegistry)
    33  //  pe.Export(w)
    34  type PrometheusExporter struct {
    35  	families map[string]*prometheusgo.MetricFamily
    36  }
    37  
    38  // MakePrometheusExporter returns an initialized prometheus exporter.
    39  func MakePrometheusExporter() PrometheusExporter {
    40  	return PrometheusExporter{families: map[string]*prometheusgo.MetricFamily{}}
    41  }
    42  
    43  // find the family for the passed-in metric, or create and return it if not found.
    44  func (pm *PrometheusExporter) findOrCreateFamily(
    45  	prom PrometheusExportable,
    46  ) *prometheusgo.MetricFamily {
    47  	familyName := exportedName(prom.GetName())
    48  	if family, ok := pm.families[familyName]; ok {
    49  		return family
    50  	}
    51  
    52  	family := &prometheusgo.MetricFamily{
    53  		Name: proto.String(familyName),
    54  		Help: proto.String(prom.GetHelp()),
    55  		Type: prom.GetType(),
    56  	}
    57  
    58  	pm.families[familyName] = family
    59  	return family
    60  }
    61  
    62  // ScrapeRegistry scrapes all metrics contained in the registry to the metric
    63  // family map, holding on only to the scraped data (which is no longer
    64  // connected to the registry and metrics within) when returning from the the
    65  // call. It creates new families as needed.
    66  func (pm *PrometheusExporter) ScrapeRegistry(registry *Registry) {
    67  	labels := registry.getLabels()
    68  	registry.Each(func(_ string, v interface{}) {
    69  		if prom, ok := v.(PrometheusExportable); ok {
    70  			m := prom.ToPrometheusMetric()
    71  			// Set registry and metric labels.
    72  			m.Label = append(labels, prom.GetLabels()...)
    73  
    74  			family := pm.findOrCreateFamily(prom)
    75  			family.Metric = append(family.Metric, m)
    76  		}
    77  	})
    78  }
    79  
    80  // PrintAsText writes all metrics in the families map to the io.Writer in
    81  // prometheus' text format. It removes individual metrics from the families
    82  // as it goes, readying the families for another found of registry additions.
    83  func (pm *PrometheusExporter) PrintAsText(w io.Writer) error {
    84  	for _, family := range pm.families {
    85  		if _, err := expfmt.MetricFamilyToText(w, family); err != nil {
    86  			return err
    87  		}
    88  	}
    89  	pm.clearMetrics()
    90  	return nil
    91  }
    92  
    93  // Verify GraphiteExporter implements Gatherer interface.
    94  var _ prometheus.Gatherer = (*PrometheusExporter)(nil)
    95  
    96  // Gather implements prometheus.Gatherer
    97  func (pm *PrometheusExporter) Gather() ([]*prometheusgo.MetricFamily, error) {
    98  	v := make([]*prometheusgo.MetricFamily, len(pm.families))
    99  	i := 0
   100  	for _, family := range pm.families {
   101  		v[i] = family
   102  		i++
   103  	}
   104  	return v, nil
   105  }
   106  
   107  // Clear metrics for reuse.
   108  func (pm *PrometheusExporter) clearMetrics() {
   109  	for _, family := range pm.families {
   110  		// Set to nil to avoid allocation if the family never gets any metrics.
   111  		family.Metric = nil
   112  	}
   113  }