From e2a71ea0ab75eb085f65b326434cd2c1b8748955 Mon Sep 17 00:00:00 2001
From: Matas Gaulia <39590298+MGaulia@users.noreply.github.com>
Date: Thu, 21 Mar 2024 17:01:17 +0100
Subject: [PATCH] add possibility to record histogram with weight (#252)

* add possibility to record histogram with weight

* add the same for recording duration
---
 stats.go      | 24 ++++++++++++++++++++----
 stats_test.go |  7 +++++++
 types.go      | 15 ++++++++++++---
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/stats.go b/stats.go
index 204c3067..8e5c7540 100644
--- a/stats.go
+++ b/stats.go
@@ -364,7 +364,7 @@ func (h *histogram) cachedReport() {
 	}
 }
 
-func (h *histogram) RecordValue(value float64) {
+func (h *histogram) recordValueWithWeight(value float64, weight int64) {
 	if h.htype != valueHistogramType {
 		return
 	}
@@ -376,10 +376,18 @@ func (h *histogram) RecordValue(value float64) {
 	idx := sort.Search(len(h.buckets), func(i int) bool {
 		return h.buckets[i].valueUpperBound >= value
 	})
-	h.samples[idx].counter.Inc(1)
+	h.samples[idx].counter.Inc(weight)
 }
 
-func (h *histogram) RecordDuration(value time.Duration) {
+func (h *histogram) RecordValue(value float64) {
+	h.recordValueWithWeight(value, 1)
+}
+
+func (h *histogram) RecordValueWithWeight(value float64, weight int64) {
+	h.recordValueWithWeight(value, weight)
+}
+
+func (h *histogram) recordDurationWithWeight(value time.Duration, weight int64) {
 	if h.htype != durationHistogramType {
 		return
 	}
@@ -391,7 +399,15 @@ func (h *histogram) RecordDuration(value time.Duration) {
 	idx := sort.Search(len(h.buckets), func(i int) bool {
 		return h.buckets[i].durationUpperBound >= value
 	})
-	h.samples[idx].counter.Inc(1)
+	h.samples[idx].counter.Inc(weight)
+}
+
+func (h *histogram) RecordDuration(value time.Duration) {
+	h.recordDurationWithWeight(value, 1)
+}
+
+func (h *histogram) RecordDurationWithWeight(value time.Duration, weight int64) {
+	h.recordDurationWithWeight(value, weight)
 }
 
 func (h *histogram) Start() Stopwatch {
diff --git a/stats_test.go b/stats_test.go
index 7d29e90e..2826087f 100644
--- a/stats_test.go
+++ b/stats_test.go
@@ -139,11 +139,14 @@ func TestHistogramValueSamples(t *testing.T) {
 	for i := 0; i < 5; i++ {
 		h.RecordValue(offset + rand.Float64()*10)
 	}
+	offset = 60
+	h.RecordValueWithWeight(offset+rand.Float64()*10, 2)
 
 	h.report(h.name, h.tags, r)
 
 	assert.Equal(t, 3, r.valueSamples[10.0])
 	assert.Equal(t, 5, r.valueSamples[60.0])
+	assert.Equal(t, 2, r.valueSamples[70.0])
 	assert.Equal(t, buckets, r.buckets)
 }
 
@@ -163,10 +166,14 @@ func TestHistogramDurationSamples(t *testing.T) {
 		h.RecordDuration(offset +
 			time.Duration(rand.Float64()*float64(10*time.Millisecond)))
 	}
+	offset = 60 * time.Millisecond
+	h.RecordDurationWithWeight(offset+
+		time.Duration(rand.Float64()*float64(10*time.Millisecond)), 2)
 
 	h.report(h.name, h.tags, r)
 
 	assert.Equal(t, 3, r.durationSamples[10*time.Millisecond])
 	assert.Equal(t, 5, r.durationSamples[60*time.Millisecond])
+	assert.Equal(t, 2, r.durationSamples[70*time.Millisecond])
 	assert.Equal(t, buckets, r.buckets)
 }
diff --git a/types.go b/types.go
index 1a15971d..3a3e98f0 100644
--- a/types.go
+++ b/types.go
@@ -30,9 +30,10 @@ import (
 // all emitted values have a given prefix or set of tags.
 //
 // IMPORTANT: When using Prometheus reporters, users must take care to
-//            not create metrics from both parent scopes and subscopes
-//            that have the same metric name but different tag keys,
-//            as metric allocation will panic.
+//
+//	not create metrics from both parent scopes and subscopes
+//	that have the same metric name but different tag keys,
+//	as metric allocation will panic.
 type Scope interface {
 	// Counter returns the Counter object corresponding to the name.
 	Counter(name string) Counter
@@ -91,10 +92,18 @@ type Histogram interface {
 	// Will use the configured value buckets for the histogram.
 	RecordValue(value float64)
 
+	// RecordValueWithWeight records a specific value directly with a weight.
+	// Will use the configured value buckets for the histogram.
+	RecordValueWithWeight(value float64, weight int64)
+
 	// RecordDuration records a specific duration directly.
 	// Will use the configured duration buckets for the histogram.
 	RecordDuration(value time.Duration)
 
+	// RecordDurationWithWeight records a specific duration directly with a weight.
+	// Will use the configured duration buckets for the histogram.
+	RecordDurationWithWeight(value time.Duration, weight int64)
+
 	// Start gives you a specific point in time to then record a duration.
 	// Will use the configured duration buckets for the histogram.
 	Start() Stopwatch