From 8496756f6dba0a23a5fde5bd986f4e3233e55ef6 Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@soundcloud.com>
Date: Thu, 10 Nov 2016 18:30:43 +0100
Subject: [PATCH] Turn timer helper into a struct with methods

---
 prometheus/timer.go      | 34 ++++++++++++++++++++++++++++------
 prometheus/timer_test.go |  4 ++--
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/prometheus/timer.go b/prometheus/timer.go
index 31df7f8..f15aba2 100644
--- a/prometheus/timer.go
+++ b/prometheus/timer.go
@@ -15,15 +15,37 @@ package prometheus
 
 import "time"
 
-type Observable interface {
+// Observer is the interface that wraps the Observe method, used by Histogram
+// and Summary to add observations.
+type Observer interface {
 	Observe(float64)
 }
 
-type TimerFunc func(Observable)
+type Timer struct {
+	begin    time.Time
+	observer Observer
+	gauge    Gauge
+}
 
-func NewTimer() TimerFunc {
-	begin := time.Now()
-	return func(obs Observable) {
-		obs.Observe(time.Since(begin).Seconds())
+func StartTimer() *Timer {
+	return &Timer{begin: time.Now()}
+}
+
+func (t *Timer) With(o Observer) *Timer {
+	t.observer = o
+	return t
+}
+
+func (t *Timer) WithGauge(g Gauge) *Timer {
+	t.gauge = g
+	return t
+}
+
+func (t *Timer) Stop() {
+	if t.observer != nil {
+		t.observer.Observe(time.Since(t.begin).Seconds())
+	}
+	if t.gauge != nil {
+		t.gauge.Set(time.Since(t.begin).Seconds())
 	}
 }
diff --git a/prometheus/timer_test.go b/prometheus/timer_test.go
index 3ca7455..747a7ba 100644
--- a/prometheus/timer_test.go
+++ b/prometheus/timer_test.go
@@ -26,8 +26,8 @@ func TestTimerObserve(t *testing.T) {
 		},
 	)
 
-	stop := NewTimer()
-	stop(his)
+	timer := StartTimer().With(his)
+	timer.Stop()
 	m := &dto.Metric{}
 	his.Write(m)