client_golang prometheus extension to push metrics to self

pushself.go

package prometheus

import (
	"fmt"
	"sort"
	"sync"

	"github.com/golang/protobuf/proto"
	dto "github.com/prometheus/client_model/go"
)

// PushMetric captures a name and metric
type PushMetric struct {
	Name   string
	Metric *dto.Metric
}

// PushSelf pushes the metrics from the default registry
func PushSelf(c chan<- PushMetric) error {
	return defRegistry.PushSelf(c)
}

// PushSelf pushes the metrics from the registry
func (r *registry) PushSelf(c chan<- PushMetric) error {

	// func (r *registry) writePB(w io.Writer, writeEncoded encoder) (int, error) {
	var metricHashes map[uint64]struct{}
	if r.collectChecksEnabled {
		metricHashes = make(map[uint64]struct{})
	}
	metricChan := make(chan Metric, capMetricChan)
	wg := sync.WaitGroup{}

	r.mtx.RLock()
	metricFamiliesByName := make(map[string]*dto.MetricFamily, len(r.dimHashesByName))

	// Scatter.
	// (Collectors could be complex and slow, so we call them all at once.)
	wg.Add(len(r.collectorsByID))
	go func() {
		wg.Wait()
		close(metricChan)
	}()
	for _, collector := range r.collectorsByID {
		go func(collector Collector) {
			defer wg.Done()
			collector.Collect(metricChan)
		}(collector)
	}
	r.mtx.RUnlock()

	// Drain metricChan in case of premature return.
	defer func() {
		for _ = range metricChan {
		}
	}()

	// Gather.
	for metric := range metricChan {
		// This could be done concurrently, too, but it required locking
		// of metricFamiliesByName (and of metricHashes if checks are
		// enabled). Most likely not worth it.
		desc := metric.Desc()
		metricFamily, ok := metricFamiliesByName[desc.fqName]
		if !ok {
			metricFamily = r.getMetricFamily()
			defer r.giveMetricFamily(metricFamily)
			metricFamily.Name = proto.String(desc.fqName)
			metricFamily.Help = proto.String(desc.help)
			metricFamiliesByName[desc.fqName] = metricFamily
		}
		dtoMetric := r.getMetric()
		defer r.giveMetric(dtoMetric)
		if err := metric.Write(dtoMetric); err != nil {
			// TODO: Consider different means of error reporting so
			// that a single erroneous metric could be skipped
			// instead of blowing up the whole collection.
			return fmt.Errorf("error collecting metric %v: %s", desc, err)
		}
		switch {
		case metricFamily.Type != nil:
			// Type already set. We are good.
			break
		case dtoMetric.Gauge != nil:
			metricFamily.Type = dto.MetricType_GAUGE.Enum()
		case dtoMetric.Counter != nil:
			metricFamily.Type = dto.MetricType_COUNTER.Enum()
		case dtoMetric.Summary != nil:
			metricFamily.Type = dto.MetricType_SUMMARY.Enum()
		case dtoMetric.Untyped != nil:
			metricFamily.Type = dto.MetricType_UNTYPED.Enum()
		case dtoMetric.Histogram != nil:
			metricFamily.Type = dto.MetricType_HISTOGRAM.Enum()
		default:
			return fmt.Errorf("empty metric collected: %s", dtoMetric)
		}
		if r.collectChecksEnabled {
			if err := r.checkConsistency(metricFamily, dtoMetric, desc, metricHashes); err != nil {
				return err
			}
		}
		metricFamily.Metric = append(metricFamily.Metric, dtoMetric)
	}

	if r.metricFamilyInjectionHook != nil {
		for _, mf := range r.metricFamilyInjectionHook() {
			existingMF, exists := metricFamiliesByName[mf.GetName()]
			if !exists {
				metricFamiliesByName[mf.GetName()] = mf
				if r.collectChecksEnabled {
					for _, m := range mf.Metric {
						if err := r.checkConsistency(mf, m, nil, metricHashes); err != nil {
							return err
						}
					}
				}
				continue
			}
			for _, m := range mf.Metric {
				if r.collectChecksEnabled {
					if err := r.checkConsistency(existingMF, m, nil, metricHashes); err != nil {
						return err
					}
				}
				existingMF.Metric = append(existingMF.Metric, m)
			}
		}
	}

	// Now that MetricFamilies are all set, sort their Metrics
	// lexicographically by their label values.
	for _, mf := range metricFamiliesByName {
		sort.Sort(metricSorter(mf.Metric))
	}

	// Write out MetricFamilies sorted by their name.
	names := make([]string, 0, len(metricFamiliesByName))
	for name := range metricFamiliesByName {
		names = append(names, name)
	}
	sort.Strings(names)

	for _, name := range names {
		for _, m := range metricFamiliesByName[name].Metric {

			c <- PushMetric{name, m}
		}
	}
	return nil
}