feat: implement prometheus histograms

cf #28

test: add basic prometheus histogram test case

Author: clementnuss

prometheus_histogram.go

@@ -0,0 +1,180 @@
+package metrics
+
+import (
+	"fmt"
+	"io"
+	"math"
+	"sync"
+)
+
+var defaultUpperBounds = []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10}
+
+// Prometheus Histogram is a histogram for non-negative values with pre-defined buckets
+//
+// Each bucket contains a counter for values in the given range.
+// Each bucket is exposed via the following metric:
+//
+//	<metric_name>_bucket{<optional_tags>,le="upper_bound"} <counter>
+//
+// Where:
+//
+//   - <metric_name> is the metric name passed to NewHistogram
+//   - <optional_tags> is optional tags for the <metric_name>, which are passed to NewHistogram
+//   - <upper_bound> - upper bound of the current bucket. all samples <= upper_bound are in that bucket
+//   - <counter> - the number of hits to the given bucket during Update* calls
+//
+// Zero histogram is usable.
+type PrometheusHistogram struct {
+	// Mu gurantees synchronous update for all the counters and sum.
+	//
+	// Do not use sync.RWMutex, since it has zero sense from performance PoV.
+	// It only complicates the code.
+	mu sync.Mutex
+
+	upperBounds []float64
+	buckets     []uint64
+
+	// count is the counter for all observations on this histogram
+	count uint64
+
+	// sum is the sum of all the values put into Histogram
+	sum float64
+}
+
+// Reset resets the given histogram.
+func (h *PrometheusHistogram) Reset() {
+	h.mu.Lock()
+	for i := range h.buckets {
+		h.buckets[i] = 0
+	}
+
+	h.sum = 0
+	h.count = 0
+	h.mu.Unlock()
+}
+
+// Update updates h with v.
+//
+// Negative values and NaNs are ignored.
+func (h *PrometheusHistogram) Update(v float64) {
+	if math.IsNaN(v) || v < 0 {
+		// Skip NaNs and negative values.
+		return
+	}
+	bucketIdx := -1
+	for i, ub := range h.upperBounds {
+		if v <= ub {
+			bucketIdx = i
+			break
+		}
+	}
+	h.mu.Lock()
+	h.sum += v
+	h.count++
+	if bucketIdx == -1 {
+		// +Inf, nothing to do, already accounted for in the total count
+	}
+	h.buckets[bucketIdx]++
+	h.mu.Unlock()
+}
+
+// Merge merges src to h
+func (h *PrometheusHistogram) Merge(src *PrometheusHistogram) {
+	// first we must compare if the upper bounds are identical
+
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	src.mu.Lock()
+	defer src.mu.Unlock()
+
+	h.sum += src.sum
+	h.count += src.count
+
+	// TODO: implement actual sum
+}
+
+// NewPrometheusHistogram creates and returns new prometheus histogram with the given name.
+//
+// name must be valid Prometheus-compatible metric with possible labels.
+// For instance,
+//
+//   - foo
+//   - foo{bar="baz"}
+//   - foo{bar="baz",aaa="b"}
+//
+// The returned histogram is safe to use from concurrent goroutines.
+func NewPrometheusHistogram(name string) *PrometheusHistogram {
+	return defaultSet.NewPrometheusHistogram(name)
+}
+
+// GetOrCreateHistogram returns registered histogram with the given name
+// or creates new histogram if the registry doesn't contain histogram with
+// the given name.
+//
+// name must be valid Prometheus-compatible metric with possible labels.
+// For instance,
+//
+//   - foo
+//   - foo{bar="baz"}
+//   - foo{bar="baz",aaa="b"}
+//
+// The returned histogram is safe to use from concurrent goroutines.
+//
+// Performance tip: prefer NewHistogram instead of GetOrCreateHistogram.
+func GetOrCreatePrometheusHistogram(name string) *PrometheusHistogram {
+	return defaultSet.GetOrCreatePrometheusHistogram(name)
+}
+
+func newPrometheusHistogram(upperBounds []float64) *PrometheusHistogram {
+	validateBuckets(&upperBounds)
+	oh := PrometheusHistogram{
+		upperBounds: upperBounds,
+		buckets:     make([]uint64, len(upperBounds)),
+	}
+
+	return &oh
+}
+
+func validateBuckets(upperBounds *[]float64) {
+	// TODO
+	if len(*upperBounds) == 0 {
+		panic("not good")
+	}
+}
+
+func (h *PrometheusHistogram) marshalTo(prefix string, w io.Writer) {
+	cumulativeSum := uint64(0)
+	h.mu.Lock()
+	count := h.count
+	sum := h.sum
+	if count == 0 {
+		h.mu.Unlock()
+		return
+	}
+	for i, ub := range h.upperBounds {
+		cumulativeSum += h.buckets[i]
+		tag := fmt.Sprintf(`le="%v"`, ub)
+		metricName := addTag(prefix, tag)
+		name, labels := splitMetricName(metricName)
+		fmt.Fprintf(w, "%s_bucket%s %d\n", name, labels, cumulativeSum)
+	}
+	h.mu.Unlock()
+
+	tag := fmt.Sprintf("le=%q", "+Inf")
+	metricName := addTag(prefix, tag)
+	name, labels := splitMetricName(metricName)
+	fmt.Fprintf(w, "%s_bucket%s %d\n", name, labels, count)
+
+	name, labels = splitMetricName(prefix)
+	if float64(int64(sum)) == sum {
+		fmt.Fprintf(w, "%s_sum%s %d\n", name, labels, int64(sum))
+	} else {
+		fmt.Fprintf(w, "%s_sum%s %g\n", name, labels, sum)
+	}
+	fmt.Fprintf(w, "%s_count%s %d\n", name, labels, count)
+}
+
+func (h *PrometheusHistogram) metricType() string {
+	return "histogram"
+}

prometheus_histogram_test.go

@@ -0,0 +1,42 @@
+package metrics
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+)
+
+func TestPrometheusHistogramSerial(t *testing.T) {
+	name := "TestPrometheusHistogramSerial"
+	h := NewPrometheusHistogram(name)
+
+	// Verify that the histogram is invisible in the output of WritePrometheus when it has no data.
+	var bb bytes.Buffer
+	WritePrometheus(&bb, false)
+	result := bb.String()
+	if strings.Contains(result, name) {
+		t.Fatalf("histogram %s shouldn't be visible in the WritePrometheus output; got\n%s", name, result)
+	}
+
+	// Write data to histogram
+	for i := 98; i < 218; i++ {
+		h.Update(float64(i)*1e-4)
+	}
+
+	// Make sure the histogram prints <prefix>_bucket on marshalTo call
+	testMarshalTo(t, h, "prefix", `prefix_bucket{le="0.005"} 0
+prefix_bucket{le="0.01"} 3
+prefix_bucket{le="0.025"} 120
+prefix_bucket{le="0.05"} 120
+prefix_bucket{le="0.1"} 120
+prefix_bucket{le="0.25"} 120
+prefix_bucket{le="0.5"} 120
+prefix_bucket{le="1"} 120
+prefix_bucket{le="2.5"} 120
+prefix_bucket{le="5"} 120
+prefix_bucket{le="10"} 120
+prefix_bucket{le="+Inf"} 120
+prefix_sum 1.8900000000000003
+prefix_count 120
+`)
+}

set.go

@@ -128,6 +128,64 @@ func (s *Set) GetOrCreateHistogram(name string) *Histogram {
 	return h
 }
 
+// NewPrometheusHistogram creates and returns new Prometheus histogram in s with the given name.
+//
+// name must be valid Prometheus-compatible metric with possible labels.
+// For instance,
+//
+//   - foo
+//   - foo{bar="baz"}
+//   - foo{bar="baz",aaa="b"}
+//
+// The returned histogram is safe to use from concurrent goroutines.
+func (s *Set) NewPrometheusHistogram(name string) *PrometheusHistogram {
+	h := newPrometheusHistogram(defaultUpperBounds)
+	s.registerMetric(name, h)
+	return h
+}
+
+// GetOrCreatePrometheusHistogram returns registered histogram in s with the given name
+// or creates new histogram if s doesn't contain histogram with the given name.
+//
+// name must be valid Prometheus-compatible metric with possible labels.
+// For instance,
+//
+//   - foo
+//   - foo{bar="baz"}
+//   - foo{bar="baz",aaa="b"}
+//
+// The returned histogram is safe to use from concurrent goroutines.
+//
+// Performance tip: prefer NewPrometheusHistogram instead of GetOrCreatePrometheusHistogram.
+func (s *Set) GetOrCreatePrometheusHistogram(name string) *PrometheusHistogram {
+	s.mu.Lock()
+	nm := s.m[name]
+	s.mu.Unlock()
+	if nm == nil {
+		// Slow path - create and register missing histogram.
+		if err := ValidateMetric(name); err != nil {
+			panic(fmt.Errorf("BUG: invalid metric name %q: %s", name, err))
+		}
+		nmNew := &namedMetric{
+			name:   name,
+			metric: newPrometheusHistogram(defaultUpperBounds),
+		}
+		s.mu.Lock()
+		nm = s.m[name]
+		if nm == nil {
+			nm = nmNew
+			s.m[name] = nm
+			s.a = append(s.a, nm)
+		}
+		s.mu.Unlock()
+	}
+	h, ok := nm.metric.(*PrometheusHistogram)
+	if !ok {
+		panic(fmt.Errorf("BUG: metric %q isn't a Histogram. It is %T", name, nm.metric))
+	}
+	return h
+}
+
 // NewCounter registers and returns new counter with the given name in the s.
 //
 // name must be valid Prometheus-compatible metric with possible labels.