From e79d7e71ce81f1f34fd755d1fb9f0b143909318f Mon Sep 17 00:00:00 2001
From: Bartlomiej Plotka <bwplotka@gmail.com>
Date: Tue, 21 Mar 2023 20:35:31 +0100
Subject: [PATCH] timer: Added support for exemplars. (#1233)

Signed-off-by: bwplotka <bwplotka@gmail.com>
---
 prometheus/timer.go      | 28 +++++++++++++++++++++-
 prometheus/timer_test.go | 51 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/prometheus/timer.go b/prometheus/timer.go
index f28a76f..52344fe 100644
--- a/prometheus/timer.go
+++ b/prometheus/timer.go
@@ -23,7 +23,9 @@ type Timer struct {
 }
 
 // NewTimer creates a new Timer. The provided Observer is used to observe a
-// duration in seconds. Timer is usually used to time a function call in the
+// duration in seconds. If the Observer implements ExemplarObserver, passing exemplar
+// later on will be also supported.
+// Timer is usually used to time a function call in the
 // following way:
 //
 //	func TimeMe() {
@@ -31,6 +33,14 @@ type Timer struct {
 //	    defer timer.ObserveDuration()
 //	    // Do actual work.
 //	}
+//
+// or
+//
+//	func TimeMeWithExemplar() {
+//		    timer := NewTimer(myHistogram)
+//		    defer timer.ObserveDurationWithExemplar(exemplar)
+//		    // Do actual work.
+//		}
 func NewTimer(o Observer) *Timer {
 	return &Timer{
 		begin:    time.Now(),
@@ -53,3 +63,19 @@ func (t *Timer) ObserveDuration() time.Duration {
 	}
 	return d
 }
+
+// ObserveDurationWithExemplar is like ObserveDuration, but it will also
+// observe exemplar with the duration unless exemplar is nil or provided Observer can't
+// be casted to ExemplarObserver.
+func (t *Timer) ObserveDurationWithExemplar(exemplar Labels) time.Duration {
+	d := time.Since(t.begin)
+	eo, ok := t.observer.(ExemplarObserver)
+	if ok && exemplar != nil {
+		eo.ObserveWithExemplar(d.Seconds(), exemplar)
+		return d
+	}
+	if t.observer != nil {
+		t.observer.Observe(d.Seconds())
+	}
+	return d
+}
diff --git a/prometheus/timer_test.go b/prometheus/timer_test.go
index 99f67db..a27912d 100644
--- a/prometheus/timer_test.go
+++ b/prometheus/timer_test.go
@@ -14,8 +14,11 @@
 package prometheus
 
 import (
+	"reflect"
 	"testing"
 
+	"google.golang.org/protobuf/proto"
+
 	dto "github.com/prometheus/client_model/go"
 )
 
@@ -52,6 +55,54 @@ func TestTimerObserve(t *testing.T) {
 	}
 }
 
+func TestTimerObserveWithExemplar(t *testing.T) {
+	var (
+		exemplar = Labels{"foo": "bar"}
+		his      = NewHistogram(HistogramOpts{Name: "test_histogram"})
+		sum      = NewSummary(SummaryOpts{Name: "test_summary"})
+		gauge    = NewGauge(GaugeOpts{Name: "test_gauge"})
+	)
+
+	func() {
+		hisTimer := NewTimer(his)
+		sumTimer := NewTimer(sum)
+		gaugeTimer := NewTimer(ObserverFunc(gauge.Set))
+		defer hisTimer.ObserveDurationWithExemplar(exemplar)
+		// Gauges and summaries does not implement ExemplarObserver, so we expect them to ignore exemplar.
+		defer sumTimer.ObserveDurationWithExemplar(exemplar)
+		defer gaugeTimer.ObserveDurationWithExemplar(exemplar)
+	}()
+
+	m := &dto.Metric{}
+	his.Write(m)
+	if want, got := uint64(1), m.GetHistogram().GetSampleCount(); want != got {
+		t.Errorf("want %d observations for histogram, got %d", want, got)
+	}
+	var got []*dto.LabelPair
+	for _, b := range m.GetHistogram().GetBucket() {
+		if b.Exemplar != nil {
+			got = b.Exemplar.GetLabel()
+			break
+		}
+	}
+
+	want := []*dto.LabelPair{{Name: proto.String("foo"), Value: proto.String("bar")}}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("expected %v exemplar labels, got %v", want, got)
+	}
+
+	m.Reset()
+	sum.Write(m)
+	if want, got := uint64(1), m.GetSummary().GetSampleCount(); want != got {
+		t.Errorf("want %d observations for summary, got %d", want, got)
+	}
+	m.Reset()
+	gauge.Write(m)
+	if got := m.GetGauge().GetValue(); got <= 0 {
+		t.Errorf("want value > 0 for gauge, got %f", got)
+	}
+}
+
 func TestTimerEmpty(t *testing.T) {
 	emptyTimer := NewTimer(nil)
 	emptyTimer.ObserveDuration()