Faster histogram implementation

This implementation is especially optimized for parallel updates, while
the decrease is negligable in serial updates.

```
                                          │ before.txt  │              after.txt              │
                                          │   sec/op    │   sec/op     vs base                │
CounterParallel/Histogram.Update_int-10     93.28n ± 2%   93.50n ± 1%        ~ (p=0.796 n=10)
CounterParallel/Histogram.Update_float-10   447.9n ± 5%   249.7n ± 7%  -44.25% (p=0.000 n=10)
CounterSerial/Histogram.Update_int-10       11.39n ± 3%   12.25n ± 2%   +7.60% (p=0.000 n=10)
CounterSerial/Histogram.Update_float-10     11.87n ± 3%   18.11n ± 0%  +52.68% (p=0.000 n=10)
geomean                                     48.74n        47.71n        -2.12%
```

Author: mattrobenolt

counter_bench_test.go

@@ -48,7 +48,33 @@ func BenchmarkCounterParallel(b *testing.B) {
 	benchmarkCounterParallel(b, "Uint64.Set", NewSet().NewUint64, (*Uint64).Set, 1)
 	benchmarkCounterParallel(b, "Float64.Add", NewSet().NewFloat64, (*Float64).Add, 1)
 	benchmarkCounterParallel(b, "Float64.Set", NewSet().NewFloat64, (*Float64).Set, 1)
-	benchmarkCounterParallel(b, "Histogram.Update", NewSet().NewHistogram, (*Histogram).Update, 1)
+	benchmarkCounterParallel(b, "Histogram.Update int", NewSet().NewHistogram, (*Histogram).Update, 1)
+	benchmarkCounterParallel(b, "Histogram.Update float", NewSet().NewHistogram, (*Histogram).Update, 1.3)
+}
+
+func BenchmarkCounterSerial(b *testing.B) {
+	benchmarkCounter(b, "Uint64.Add", NewSet().NewUint64, (*Uint64).Add, 1)
+	benchmarkCounter(b, "Uint64.Set", NewSet().NewUint64, (*Uint64).Set, 1)
+	benchmarkCounter(b, "Float64.Add", NewSet().NewFloat64, (*Float64).Add, 1)
+	benchmarkCounter(b, "Float64.Set", NewSet().NewFloat64, (*Float64).Set, 1)
+	benchmarkCounter(b, "Histogram.Update int", NewSet().NewHistogram, (*Histogram).Update, 1)
+	benchmarkCounter(b, "Histogram.Update float", NewSet().NewHistogram, (*Histogram).Update, 1.3)
+}
+
+func benchmarkCounter[T any, V any](
+	b *testing.B,
+	name string,
+	setup func(string, ...string) *T,
+	do func(*T, V),
+	value V,
+) {
+	b.Helper()
+	thing := setup("foo")
+	b.Run(name, func(b *testing.B) {
+		for b.Loop() {
+			do(thing, value)
+		}
+	})
 }
 
 func benchmarkCounterParallel[T any, V any](
@@ -61,7 +87,6 @@ func benchmarkCounterParallel[T any, V any](
 	b.Helper()
 	thing := setup("foo")
 	b.Run(name, func(b *testing.B) {
-		b.ReportAllocs()
 		b.RunParallel(func(pb *testing.PB) {
 			for pb.Next() {
 				do(thing, value)

histogram.go

@@ -104,17 +104,15 @@ type Histogram struct {
 	upper atomic.Uint64
 
 	// sum is the sum of all the values put into Histogram
-	sumInt   atomic.Uint64
-	sumFloat atomicx.Float64
+	sum atomicx.Sum
 }
 
 // Reset resets the given histogram.
 func (h *Histogram) Reset() {
 	clear(h.buckets[:])
 	h.lower.Store(0)
 	h.upper.Store(0)
-	h.sumInt.Store(0)
-	h.sumFloat.Store(0)
+	h.sum.Reset()
 }
 
 // Update updates h with val.
@@ -157,13 +155,7 @@ func (h *Histogram) Update(val float64) {
 		db[offset].Add(1)
 	}
 
-	if val > 0 {
-		if intval := uint64(val); float64(intval) == val {
-			h.sumInt.Add(intval)
-		} else {
-			h.sumFloat.Add(val)
-		}
-	}
+	h.sum.Add(val)
 }
 
 // Observe updates h with val, identical to [Histogram.Update].
@@ -177,8 +169,7 @@ func (h *Histogram) Observe(val float64) {
 func (h *Histogram) Merge(src *Histogram) {
 	h.lower.Add(src.lower.Load())
 	h.upper.Add(src.upper.Load())
-	h.sumInt.Add(src.sumInt.Load())
-	h.sumFloat.Add(src.sumFloat.Load())
+	h.sum.Add(src.sum.Load())
 
 	for i := range src.buckets {
 		if dbSrc := src.buckets[i].Load(); dbSrc != nil {
@@ -216,7 +207,7 @@ func (h *Histogram) marshalTo(w ExpfmtWriter, name MetricName) {
 		return
 	}
 
-	sum := h.sum()
+	sum := h.sum.Load()
 	family := name.Family.String()
 
 	// 1 extra because we're always adding in the vmrange tag
@@ -329,10 +320,6 @@ func (h *Histogram) punchBuckets(c *punchCard) (total uint64, punches int) {
 	return
 }
 
-func (h *Histogram) sum() float64 {
-	return float64(h.sumInt.Load()) + h.sumFloat.Load()
-}
-
 // punchCard is used internally to track counts per bucket when computing
 // which histograms ranges have been hit.
 type punchCard [totalBuckets]uint64

internal/atomicx/atomicx.go

@@ -2,16 +2,30 @@ package atomicx
 
 import (
 	"math"
+	"math/rand/v2"
 	"sync/atomic"
 
 	"go.withmatt.com/metrics/internal/fasttime"
 )
 
-// An Float64 is an atomic float64. The zero value is zero.
+// A Float64 is an atomic float64. The zero value is zero.
 type Float64 struct {
 	v atomic.Uint64
 }
 
+func (x *Float64) Add(val float64) {
+	if val == 0 {
+		return
+	}
+	for {
+		oldBits := x.v.Load()
+		newBits := math.Float64bits(math.Float64frombits(oldBits) + val)
+		if x.v.CompareAndSwap(oldBits, newBits) {
+			return
+		}
+	}
+}
+
 func (x *Float64) Load() float64 {
 	return math.Float64frombits(x.v.Load())
 }
@@ -20,19 +34,51 @@ func (x *Float64) Store(val float64) {
 	x.v.Store(math.Float64bits(val))
 }
 
-func (x *Float64) Add(val float64) {
+// A Sum is an atomic float64 that can only Add. The zero value is zero.
+type Sum struct {
+	i atomic.Uint64
+	v [2]atomic.Uint64
+}
+
+func (x *Sum) AddUint64(val uint64) {
 	if val == 0 {
 		return
 	}
+	x.i.Add(val)
+}
+
+func (x *Sum) Add(val float64) {
+	if val <= 0 {
+		return
+	}
+
+	if intval := uint64(val); val == float64(intval) {
+		x.i.Add(intval)
+		return
+	}
+
+	//nolint:gosec
+	idx := rand.Uint64() & 1
 	for {
-		oldBits := x.v.Load()
+		oldBits := x.v[idx].Load()
 		newBits := math.Float64bits(math.Float64frombits(oldBits) + val)
-		if x.v.CompareAndSwap(oldBits, newBits) {
+		if x.v[idx].CompareAndSwap(oldBits, newBits) {
 			return
 		}
 	}
 }
 
+func (x *Sum) Reset() {
+	x.i.Store(0)
+	clear(x.v[:])
+}
+
+func (x *Sum) Load() float64 {
+	return float64(x.i.Load()) +
+		math.Float64frombits(x.v[0].Load()) +
+		math.Float64frombits(x.v[1].Load())
+}
+
 // An Instant is an atomic fasttime.Instant. The zero value is zero.
 type Instant struct {
 	v atomic.Int64

internal/atomicx/atomicx_bench_test.go

@@ -0,0 +1,91 @@
+package atomicx
+
+import "testing"
+
+func BenchmarkLoadParallel(b *testing.B) {
+	b.Run("Float64", func(b *testing.B) {
+		var f Float64
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				f.Load()
+			}
+		})
+	})
+
+	b.Run("Sum", func(b *testing.B) {
+		var s Sum
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				s.Load()
+			}
+		})
+	})
+}
+
+func BenchmarkAddParallel(b *testing.B) {
+	b.Run("float/type=Float64", func(b *testing.B) {
+		var f Float64
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				f.Add(1.1)
+			}
+		})
+	})
+
+	b.Run("integer/type=Float64", func(b *testing.B) {
+		var f Float64
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				f.Add(1)
+			}
+		})
+	})
+
+	b.Run("float/type=Sum", func(b *testing.B) {
+		var s Sum
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				s.Add(1.1)
+			}
+		})
+	})
+
+	b.Run("integer/type=Sum", func(b *testing.B) {
+		var s Sum
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				s.Add(1)
+			}
+		})
+	})
+}
+
+func BenchmarkAdd(b *testing.B) {
+	b.Run("float/type=Float64", func(b *testing.B) {
+		var f Float64
+		for b.Loop() {
+			f.Add(1.1)
+		}
+	})
+
+	b.Run("integer/type=Float64", func(b *testing.B) {
+		var f Float64
+		for b.Loop() {
+			f.Add(1)
+		}
+	})
+
+	b.Run("float/type=Sum", func(b *testing.B) {
+		var s Sum
+		for b.Loop() {
+			s.Add(1.1)
+		}
+	})
+
+	b.Run("integer/type=Sum", func(b *testing.B) {
+		var s Sum
+		for b.Loop() {
+			s.Add(1)
+		}
+	})
+}

internal/atomicx/atomicx_test.go

@@ -1,6 +1,7 @@
 package atomicx_test
 
 import (
+	"math"
 	"runtime"
 	"testing"
 
@@ -11,14 +12,25 @@ import (
 func TestFloat64(t *testing.T) {
 	var f Float64
 	v := float64(1.88)
-	assert.Equal(t, 0, f.Load())
+	assert.Equal(t, f.Load(), 0)
 	f.Store(v)
-	assert.Equal(t, v, f.Load())
+	assert.Equal(t, f.Load(), v)
 	f.Add(v)
-	assert.Equal(t, v+v, f.Load())
+	assert.Equal(t, f.Load(), v+v)
 }
 
-func TestHammerAdd(t *testing.T) {
+func TestSum(t *testing.T) {
+	var s Sum
+	assert.Equal(t, s.Load(), 0)
+	s.Add(1.88)
+	assert.Equal(t, s.Load(), 1.88)
+	s.Add(1)
+	assert.Equal(t, s.Load(), 2.88)
+	s.Reset()
+	assert.Equal(t, s.Load(), 0)
+}
+
+func TestHammerFloatAdd(t *testing.T) {
 	const p = 4
 	n := 100000
 	if testing.Short() {
@@ -44,3 +56,36 @@ func TestHammerAdd(t *testing.T) {
 	}
 	assert.Equal(t, val.Load(), float64(n)*p)
 }
+
+func TestHammerSumAdd(t *testing.T) {
+	const p = 4
+	n := 10000
+	if testing.Short() {
+		n = 1000
+	}
+	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(p))
+
+	const a = 1.1
+	const b = 1
+
+	c := make(chan int)
+	var val Sum
+	for range p {
+		go func() {
+			defer func() {
+				assert.Nil(t, recover())
+				c <- 1
+			}()
+			for range n {
+				val.Add(a)
+				val.Add(b)
+				val.AddUint64(b)
+			}
+		}()
+	}
+	for range p {
+		<-c
+	}
+	// XXX: Floating point precision, so need to round
+	assert.Equal(t, math.Round(val.Load()), math.Round(float64(n)*(p*a+p*b*2)))
+}

internal/atomicx/benchmarks.txt

@@ -0,0 +1,12 @@
+$ benchstat -col /type compare.txt
+goos: darwin
+goarch: arm64
+pkg: go.withmatt.com/metrics/internal/atomicx
+cpu: Apple M1 Max
+                       │   Float64    │                 Sum                 │
+                       │    sec/op    │   sec/op     vs base                │
+AddParallel/float-10      336.3n ± 7%   270.2n ± 3%  -19.64% (p=0.000 n=10)
+AddParallel/integer-10   331.75n ± 5%   80.38n ± 2%  -75.77% (p=0.000 n=10)
+Add/float-10             10.180n ± 1%   9.612n ± 1%   -5.58% (p=0.000 n=10)
+Add/integer-10           10.075n ± 1%   6.883n ± 0%  -31.69% (p=0.000 n=10)
+geomean                   58.16n        34.62n       -40.47%