From b23fdf5bd7b61dbf896560607cea350caa019e66 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Wed, 20 Dec 2023 14:06:44 +0200
Subject: [PATCH] gauge.go: add Set() method, which can be used for changing
 the gauge value without the need to pass callback to NewGauge()

---
 gauge.go      | 31 +++++++++++++++++++++++++------
 gauge_test.go | 22 ++++++++++++++++++----
 set.go        |  6 ------
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/gauge.go b/gauge.go
index d40b730..9f676f4 100644
--- a/gauge.go
+++ b/gauge.go
@@ -3,10 +3,11 @@ package metrics
 import (
 	"fmt"
 	"io"
+	"math"
+	"sync/atomic"
 )
 
-// NewGauge registers and returns gauge with the given name, which calls f
-// to obtain gauge value.
+// NewGauge registers and returns gauge with the given name, which calls f to obtain gauge value.
 //
 // name must be valid Prometheus-compatible metric with possible labels.
 // For instance,
@@ -16,6 +17,7 @@ import (
 //   - foo{bar="baz",aaa="b"}
 //
 // f must be safe for concurrent calls.
+// if f is nil, then it is expected that the gauge value is changed via Gauge.Set() call.
 //
 // The returned gauge is safe to use from concurrent goroutines.
 //
@@ -25,19 +27,36 @@ func NewGauge(name string, f func() float64) *Gauge {
 }
 
 // Gauge is a float64 gauge.
-//
-// See also Counter, which could be used as a gauge with Set and Dec calls.
 type Gauge struct {
+	// f is a callback, which is called for returning the gauge value.
 	f func() float64
+
+	// valueBits contains uint64 representation of float64 passed to Gauge.Set.
+	valueBits uint64
 }
 
 // Get returns the current value for g.
 func (g *Gauge) Get() float64 {
-	return g.f()
+	if f := g.f; f != nil {
+		return f()
+	}
+	n := atomic.LoadUint64(&g.valueBits)
+	return math.Float64frombits(n)
+}
+
+// Set sets g value to v.
+//
+// The g must be created with nil callback in order to be able to call this function.
+func (g *Gauge) Set(v float64) {
+	if g.f != nil {
+		panic(fmt.Errorf("cannot call Set on gauge created with non-nil callback"))
+	}
+	n := math.Float64bits(v)
+	atomic.StoreUint64(&g.valueBits, n)
 }
 
 func (g *Gauge) marshalTo(prefix string, w io.Writer) {
-	v := g.f()
+	v := g.Get()
 	if float64(int64(v)) == v {
 		// Marshal integer values without scientific notation
 		fmt.Fprintf(w, "%s %d\n", prefix, int64(v))
diff --git a/gauge_test.go b/gauge_test.go
index 3f6e5cd..484a4c9 100644
--- a/gauge_test.go
+++ b/gauge_test.go
@@ -7,14 +7,28 @@ import (
 )
 
 func TestGaugeError(t *testing.T) {
-	expectPanic(t, "NewGauge_nil_callback", func() {
-		NewGauge("NewGauge_nil_callback", nil)
+	expectPanic(t, "NewGauge_Set_non-nil-callback", func() {
+		g := NewGauge("NewGauge_non_nil_callback", func() float64 { return 123 })
+		g.Set(12.35)
 	})
-	expectPanic(t, "GetOrCreateGauge_nil_callback", func() {
-		GetOrCreateGauge("GetOrCreateGauge_nil_callback", nil)
+	expectPanic(t, "GetOrCreateGauge_Set_non-nil-callback", func() {
+		g := GetOrCreateGauge("GetOrCreateGauge_nil_callback", func() float64 { return 123 })
+		g.Set(42)
 	})
 }
 
+func TestGaugeSet(t *testing.T) {
+	s := NewSet()
+	g := s.NewGauge("foo", nil)
+	if n := g.Get(); n != 0 {
+		t.Fatalf("unexpected gauge value: %g; expecting 0", n)
+	}
+	g.Set(1.234)
+	if n := g.Get(); n != 1.234 {
+		t.Fatalf("unexpected gauge value %g; expecting 1.234", n)
+	}
+}
+
 func TestGaugeSerial(t *testing.T) {
 	name := "GaugeSerial"
 	n := 1.23
diff --git a/set.go b/set.go
index 9949b7c..50a095b 100644
--- a/set.go
+++ b/set.go
@@ -251,9 +251,6 @@ func (s *Set) GetOrCreateFloatCounter(name string) *FloatCounter {
 //
 // The returned gauge is safe to use from concurrent goroutines.
 func (s *Set) NewGauge(name string, f func() float64) *Gauge {
-	if f == nil {
-		panic(fmt.Errorf("BUG: f cannot be nil"))
-	}
 	g := &Gauge{
 		f: f,
 	}
@@ -280,9 +277,6 @@ func (s *Set) GetOrCreateGauge(name string, f func() float64) *Gauge {
 	s.mu.Unlock()
 	if nm == nil {
 		// Slow path - create and register missing gauge.
-		if f == nil {
-			panic(fmt.Errorf("BUG: f cannot be nil"))
-		}
 		if err := validateMetric(name); err != nil {
 			panic(fmt.Errorf("BUG: invalid metric name %q: %s", name, err))
 		}