testing.go

package ph

import (
	"sync"

	"github.com/prometheus/client_golang/prometheus"
)

type TrackingRegistry struct {
	prometheus.Registerer
	prometheus.Gatherer

	mu    sync.Mutex
	descs []*prometheus.Desc
}

func NewTrackingRegistry() *TrackingRegistry {
	reg := prometheus.NewRegistry()

	return &TrackingRegistry{
		Registerer: reg,
		Gatherer:   reg,
	}
}

func (r *TrackingRegistry) Register(cs prometheus.Collector) error {
	r.track(cs)
	return r.Registerer.(prometheus.Registerer).Register(cs)
}

func (r *TrackingRegistry) MustRegister(cs ...prometheus.Collector) {
	for _, c := range cs {
		r.track(c)
	}
	r.Registerer.(prometheus.Registerer).MustRegister(cs...)
}

func (r *TrackingRegistry) track(c prometheus.Collector) {
	ch := make(chan *prometheus.Desc, 10)

	go func() {
		c.Describe(ch)
		close(ch)
	}()

	for d := range ch {
		r.mu.Lock()
		r.descs = append(r.descs, d)
		r.mu.Unlock()
	}
}

func (r *TrackingRegistry) MetricNames() []string {
	r.mu.Lock()
	defer r.mu.Unlock()

	names := make([]string, len(r.descs))
	for i, d := range r.descs {
		names[i] = d.String()
	}

	return names
}