feat: use time-unbiased algorithm for histogram reservoir

Author: dashpole

sdk/metric/exemplar/benchmark_test.go

@@ -23,7 +23,7 @@ func BenchmarkFixedSizeReservoirOffer(b *testing.B) {
 			// number of collect calls.
 			if i%100 == 99 {
 				reservoir.mu.Lock()
-				reservoir.reset()
+				reservoir.nt.reset()
 				reservoir.mu.Unlock()
 			}
 			i++

sdk/metric/exemplar/fixed_size_reservoir.go

@@ -5,8 +5,6 @@ package exemplar // import "go.opentelemetry.io/otel/sdk/metric/exemplar"
 
 import (
 	"context"
-	"math"
-	"math/rand/v2"
 	"sync"
 	"time"
 
@@ -26,7 +24,10 @@ func FixedSizeReservoirProvider(k int) ReservoirProvider {
 // sample each one. If there are more than k, the Reservoir will then randomly
 // sample all additional measurement with a decreasing probability.
 func NewFixedSizeReservoir(k int) *FixedSizeReservoir {
-	return newFixedSizeReservoir(newStorage(k))
+	return &FixedSizeReservoir{
+		nt:      newNextTracker(k),
+		storage: make([]measurement, k),
+	}
 }
 
 var _ Reservoir = &FixedSizeReservoir{}
@@ -37,41 +38,9 @@ var _ Reservoir = &FixedSizeReservoir{}
 // additional measurement with a decreasing probability.
 type FixedSizeReservoir struct {
 	reservoir.ConcurrentSafe
-	*storage
-	mu sync.Mutex
-
-	// count is the number of measurement seen.
-	count int64
-	// next is the next count that will store a measurement at a random index
-	// once the reservoir has been filled.
-	next int64
-	// w is the largest random number in a distribution that is used to compute
-	// the next next.
-	w float64
-}
-
-func newFixedSizeReservoir(s *storage) *FixedSizeReservoir {
-	r := &FixedSizeReservoir{
-		storage: s,
-	}
-	r.reset()
-	return r
-}
-
-// randomFloat64 returns, as a float64, a uniform pseudo-random number in the
-// open interval (0.0,1.0).
-func (*FixedSizeReservoir) randomFloat64() float64 {
-	// TODO: Use an algorithm that avoids rejection sampling. For example:
-	//
-	//   const precision = 1 << 53 // 2^53
-	//   // Generate an integer in [1, 2^53 - 1]
-	//   v := rand.Uint64() % (precision - 1) + 1
-	//   return float64(v) / float64(precision)
-	f := rand.Float64()
-	for f == 0 {
-		f = rand.Float64()
-	}
-	return f
+	mu      sync.Mutex
+	storage []measurement
+	nt      *nextTracker
 }
 
 // Offer accepts the parameters associated with a measurement. The
@@ -86,108 +55,12 @@ func (*FixedSizeReservoir) randomFloat64() float64 {
 // parameters are the value and dropped (filtered) attributes of the
 // measurement respectively.
 func (r *FixedSizeReservoir) Offer(ctx context.Context, t time.Time, n Value, a []attribute.KeyValue) {
-	// The following algorithm is "Algorithm L" from Li, Kim-Hung (4 December
-	// 1994). "Reservoir-Sampling Algorithms of Time Complexity
-	// O(n(1+log(N/n)))". ACM Transactions on Mathematical Software. 20 (4):
-	// 481–493 (https://dl.acm.org/doi/10.1145/198429.198435).
-	//
-	// A high-level overview of "Algorithm L":
-	//   0) Pre-calculate the random count greater than the storage size when
-	//      an exemplar will be replaced.
-	//   1) Accept all measurements offered until the configured storage size is
-	//      reached.
-	//   2) Loop:
-	//      a) When the pre-calculate count is reached, replace a random
-	//         existing exemplar with the offered measurement.
-	//      b) Calculate the next random count greater than the existing one
-	//         which will replace another exemplars
-	//
-	// The way a "replacement" count is computed is by looking at `n` number of
-	// independent random numbers each corresponding to an offered measurement.
-	// Of these numbers the smallest `k` (the same size as the storage
-	// capacity) of them are kept as a subset. The maximum value in this
-	// subset, called `w` is used to weight another random number generation
-	// for the next count that will be considered.
-	//
-	// By weighting the next count computation like described, it is able to
-	// perform a uniformly-weighted sampling algorithm based on the number of
-	// samples the reservoir has seen so far. The sampling will "slow down" as
-	// more and more samples are offered so as to reduce a bias towards those
-	// offered just prior to the end of the collection.
-	//
-	// This algorithm is preferred because of its balance of simplicity and
-	// performance. It will compute three random numbers (the bulk of
-	// computation time) for each item that becomes part of the reservoir, but
-	// it does not spend any time on items that do not. In particular it has an
-	// asymptotic runtime of O(k(1 + log(n/k)) where n is the number of
-	// measurements offered and k is the reservoir size.
-	//
-	// See https://en.wikipedia.org/wiki/Reservoir_sampling for an overview of
-	// this and other reservoir sampling algorithms. See
-	// https://github.com/MrAlias/reservoir-sampling for a performance
-	// comparison of reservoir sampling algorithms.
-
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	if int(r.count) < cap(r.measurements) {
-		r.store(ctx, int(r.count), t, n, a)
-	} else if r.count == r.next {
-		// Overwrite a random existing measurement with the one offered.
-		idx := int(rand.Int64N(int64(cap(r.measurements))))
-		r.store(ctx, idx, t, n, a)
-		r.advance()
+	sampled, idx := r.nt.shouldSample()
+	if sampled {
+		r.storage[idx].store(ctx, t, n, a)
 	}
-	r.count++
-}
-
-// reset resets r to the initial state.
-func (r *FixedSizeReservoir) reset() {
-	// This resets the number of exemplars known.
-	r.count = 0
-	// Random index inserts should only happen after the storage is full.
-	r.next = int64(cap(r.measurements))
-
-	// Initial random number in the series used to generate r.next.
-	//
-	// This is set before r.advance to reset or initialize the random number
-	// series. Without doing so it would always be 0 or never restart a new
-	// random number series.
-	//
-	// This maps the uniform random number in (0,1) to a geometric distribution
-	// over the same interval. The mean of the distribution is inversely
-	// proportional to the storage capacity.
-	r.w = math.Exp(math.Log(r.randomFloat64()) / float64(cap(r.measurements)))
-
-	r.advance()
-}
-
-// advance updates the count at which the offered measurement will overwrite an
-// existing exemplar.
-func (r *FixedSizeReservoir) advance() {
-	// Calculate the next value in the random number series.
-	//
-	// The current value of r.w is based on the max of a distribution of random
-	// numbers (i.e. `w = max(u_1,u_2,...,u_k)` for `k` equal to the capacity
-	// of the storage and each `u` in the interval (0,w)). To calculate the
-	// next r.w we use the fact that when the next exemplar is selected to be
-	// included in the storage an existing one will be dropped, and the
-	// corresponding random number in the set used to calculate r.w will also
-	// be replaced. The replacement random number will also be within (0,w),
-	// therefore the next r.w will be based on the same distribution (i.e.
-	// `max(u_1,u_2,...,u_k)`). Therefore, we can sample the next r.w by
-	// computing the next random number `u` and take r.w as `w * u^(1/k)`.
-	r.w *= math.Exp(math.Log(r.randomFloat64()) / float64(cap(r.measurements)))
-	// Use the new random number in the series to calculate the count of the
-	// next measurement that will be stored.
-	//
-	// Given 0 < r.w < 1, each iteration will result in subsequent r.w being
-	// smaller. This translates here into the next next being selected against
-	// a distribution with a higher mean (i.e. the expected value will increase
-	// and replacements become less likely)
-	//
-	// Important to note, the new r.next will always be at least 1 more than
-	// the last r.next.
-	r.next += int64(math.Log(r.randomFloat64())/math.Log(1-r.w)) + 1
 }
 
 // Collect returns all the held exemplars.
@@ -196,10 +69,17 @@ func (r *FixedSizeReservoir) advance() {
 func (r *FixedSizeReservoir) Collect(dest *[]Exemplar) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.storage.Collect(dest)
+	*dest = reset(*dest, len(r.storage), len(r.storage))
+	var n int
+	for i := range r.storage {
+		if r.storage[i].exemplar(&(*dest)[n]) {
+			n++
+		}
+	}
+	*dest = (*dest)[:n]
 	// Call reset here even though it will reset r.count and restart the random
 	// number series. This will persist any old exemplars as long as no new
 	// measurements are offered, but it will also prioritize those new
 	// measurements that are made over the older collection cycle ones.
-	r.reset()
+	r.nt.reset()
 }

sdk/metric/exemplar/fixed_size_reservoir_test.go

@@ -45,8 +45,8 @@ func TestNewFixedSizeReservoirSamplingCorrectness(t *testing.T) {
 	}
 
 	var sum float64
-	for i := range r.measurements {
-		sum += r.measurements[i].Value.Float64()
+	for i := range r.storage {
+		sum += r.storage[i].Value.Float64()
 	}
 	mean := sum / float64(sampleSize)
 

sdk/metric/exemplar/histogram_reservoir.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"slices"
 	"sort"
+	"sync"
 	"time"
 
 	"go.opentelemetry.io/otel/attribute"
@@ -22,29 +23,38 @@ func HistogramReservoirProvider(bounds []float64) ReservoirProvider {
 	}
 }
 
-// NewHistogramReservoir returns a [HistogramReservoir] that samples the last
-// measurement that falls within a histogram bucket. The histogram bucket
-// upper-boundaries are define by bounds.
+type bucket struct {
+	mu sync.Mutex
+	nt nextTracker
+	measurement
+}
+
+// NewHistogramReservoir returns a [HistogramReservoir] that samples
+// measurements that fall within a histogram bucket using Algorithm L. The
+// histogram bucket upper-boundaries are defined by bounds.
 //
 // The passed bounds must be sorted before calling this function.
 func NewHistogramReservoir(bounds []float64) *HistogramReservoir {
+	buckets := make([]bucket, len(bounds)+1)
+	for i := range buckets {
+		buckets[i].nt = *newNextTracker(1)
+	}
 	return &HistogramReservoir{
 		bounds:  bounds,
-		storage: newStorage(len(bounds) + 1),
+		buckets: buckets,
 	}
 }
 
 var _ Reservoir = &HistogramReservoir{}
 
-// HistogramReservoir is a [Reservoir] that samples the last measurement that
-// falls within a histogram bucket. The histogram bucket upper-boundaries are
-// define by bounds.
+// HistogramReservoir is a [Reservoir] that samples
+// measurements that fall within a histogram bucket using Algorithm L. The
+// histogram bucket upper-boundaries are defined by bounds.
 type HistogramReservoir struct {
 	reservoir.ConcurrentSafe
-	*storage
-
 	// bounds are bucket bounds in ascending order.
-	bounds []float64
+	bounds  []float64
+	buckets []bucket
 }
 
 // Offer accepts the parameters associated with a measurement. The
@@ -69,14 +79,31 @@ func (r *HistogramReservoir) Offer(ctx context.Context, t time.Time, v Value, a
 		panic("unknown value type")
 	}
 
-	idx := sort.SearchFloat64s(r.bounds, n)
+	b := &r.buckets[sort.SearchFloat64s(r.bounds, n)]
 
-	r.store(ctx, idx, t, v, a)
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	sampled, _ := b.nt.shouldSample()
+	if sampled {
+		b.store(ctx, t, v, a)
+	}
 }
 
 // Collect returns all the held exemplars.
 //
 // The Reservoir state is preserved after this call.
 func (r *HistogramReservoir) Collect(dest *[]Exemplar) {
-	r.storage.Collect(dest)
+	*dest = reset(*dest, len(r.buckets), len(r.buckets))
+	var n int
+	for i := range r.buckets {
+		b := &r.buckets[i]
+		b.mu.Lock()
+		if b.exemplar(&(*dest)[n]) {
+			n++
+		}
+		b.nt.reset()
+		b.mu.Unlock()
+	}
+	*dest = (*dest)[:n]
 }

sdk/metric/exemplar/histogram_reservoir_test.go

@@ -3,7 +3,12 @@
 
 package exemplar
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
 
 func TestHist(t *testing.T) {
 	bounds := []float64{0, 100}
@@ -25,3 +30,35 @@ func TestHistogramReservoirConcurrentSafe(t *testing.T) {
 		return HistogramReservoirProvider(bounds), len(bounds)
 	}))
 }
+
+func TestHistogramReservoirTimeUnbiased(t *testing.T) {
+	bounds := []float64{10}
+	r := NewHistogramReservoir(bounds)
+
+	const (
+		N = 100   // Items per run
+		M = 20000 // Number of runs
+	)
+
+	var sum float64
+	var dest []Exemplar
+
+	for range M {
+		for j := 1; j <= N; j++ {
+			val := 10.0 * float64(j) / float64(N)
+			r.Offer(t.Context(), staticTime, NewValue(val), nil)
+		}
+		r.Collect(&dest)
+		require.Len(t, dest, 1)
+		sum += dest[0].Value.Float64()
+		dest = dest[:0]
+	}
+
+	mean := sum / float64(M)
+	expectedMean := 5.0 * float64(N+1) / float64(N)
+
+	// Standard deviation of the discrete uniform distribution is approx 2.88.
+	// Standard error of the mean for M=20000 is 2.88 / sqrt(20000) approx 0.02.
+	// A delta of 0.1 is approx 5 standard errors, which makes flakiness extremely unlikely.
+	assert.InDelta(t, expectedMean, mean, 0.1)
+}