From 3bbb4b24a47f595e62cfbae62205f641ee921f14 Mon Sep 17 00:00:00 2001
From: Raphael Simon <simon.raphael@gmail.com>
Date: Fri, 26 Jan 2024 17:50:56 -0800
Subject: [PATCH] Fix/trace sampler (#3467)

* Revert "Use crypto/rand for sampling (#3466)"

This reverts commit a17e6876928c38945a28fcded2142b011f156536.

* Use global random generator which is concurrency-safe
---
 middleware/sampler.go      | 44 ++++++++++----------------------------
 middleware/sampler_test.go | 39 ++++++++++++++++-----------------
 2 files changed, 30 insertions(+), 53 deletions(-)

diff --git a/middleware/sampler.go b/middleware/sampler.go
index 2f2f722b34..a198ddedcd 100644
--- a/middleware/sampler.go
+++ b/middleware/sampler.go
@@ -1,8 +1,7 @@
 package middleware
 
 import (
-	"crypto/rand"
-	"math/big"
+	"math/rand"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -15,16 +14,6 @@ type (
 		Sample() bool
 	}
 
-	// randomizer is an interface for generating random numbers.
-	randomizer interface {
-		// Int64 returns a random int64 in the range [0, bound).
-		Int64(bound int64) int64
-	}
-
-	// randomGenerator is a randomizer that uses crypto/rand.
-	randomGenerator struct {
-	}
-
 	adaptiveSampler struct {
 		sync.Mutex
 		lastRate        int64
@@ -37,14 +26,15 @@ type (
 	fixedSampler int
 )
 
+// Let tests override the random number generator.
+var intn = rand.Intn
+
 const (
 	// adaptive upper bound has granularity in case caller becomes extremely busy.
-	adaptiveUpperBoundInt   = int64(10000)
+	adaptiveUpperBoundInt   = 10000
 	adaptiveUpperBoundFloat = float64(adaptiveUpperBoundInt)
 )
 
-var rnd randomizer = newRandomizer()
-
 // NewAdaptiveSampler computes the interval for sampling for tracing middleware.
 // it can also be used by non-web go routines to trace internal API calls.
 //
@@ -79,14 +69,14 @@ func NewFixedSampler(samplingPercent int) Sampler {
 // Sample implementation for adaptive rate
 func (s *adaptiveSampler) Sample() bool {
 	// adjust sampling rate whenever sample size is reached.
-	var currentRate int64
+	var currentRate int
 	if atomic.AddUint32(&s.counter, 1) == s.sampleSize { // exact match prevents
 		atomic.StoreUint32(&s.counter, 0) // race is ok
 		s.Lock()
 		{
 			d := time.Since(s.start).Seconds()
 			r := float64(s.sampleSize) / d
-			currentRate = int64((float64(s.maxSamplingRate) * adaptiveUpperBoundFloat) / r)
+			currentRate = int((float64(s.maxSamplingRate) * adaptiveUpperBoundFloat) / r)
 			if currentRate > adaptiveUpperBoundInt {
 				currentRate = adaptiveUpperBoundInt
 			} else if currentRate < 1 {
@@ -97,27 +87,15 @@ func (s *adaptiveSampler) Sample() bool {
 		s.Unlock()
 		atomic.StoreInt64(&s.lastRate, int64(currentRate))
 	} else {
-		currentRate = int64(atomic.LoadInt64(&s.lastRate))
+		currentRate = int(atomic.LoadInt64(&s.lastRate))
 	}
 
 	// currentRate is never zero.
-	return currentRate == adaptiveUpperBoundInt || rnd.Int64(adaptiveUpperBoundInt) < currentRate
+	return currentRate == adaptiveUpperBoundInt || intn(adaptiveUpperBoundInt) < currentRate
 }
 
 // Sample implementation for fixed percentage
 func (s fixedSampler) Sample() bool {
-	samplingPercent := int64(s)
-	return samplingPercent > 0 && (samplingPercent == 100 || rnd.Int64(100) < samplingPercent)
-}
-
-// newRandomizer returns a randomizer.
-func newRandomizer() randomizer {
-	return &randomGenerator{}
-}
-
-// Int64 returns a random int64 in the range [0, bound).
-func (r *randomGenerator) Int64(bound int64) int64 {
-	// can we ignore the error?
-	rnd, _ := rand.Int(rand.Reader, big.NewInt(bound))
-	return rnd.Int64()
+	samplingPercent := int(s)
+	return samplingPercent > 0 && (samplingPercent == 100 || intn(100) < samplingPercent)
 }
diff --git a/middleware/sampler_test.go b/middleware/sampler_test.go
index ef2abb73a9..ee911095f1 100644
--- a/middleware/sampler_test.go
+++ b/middleware/sampler_test.go
@@ -6,24 +6,6 @@ import (
 	"time"
 )
 
-type (
-	deterministicGenerator struct {
-		*rand.Rand
-	}
-)
-
-func newDeterministicGenerator() randomizer {
-	r := rand.New(rand.NewSource(1))
-	r.Seed(123) // make the random generator deterministic
-	return &deterministicGenerator{
-		Rand: r,
-	}
-}
-
-func (d *deterministicGenerator) Int64(bound int64) int64 {
-	return int64(d.Intn(int(bound)))
-}
-
 func TestFixedSampler(t *testing.T) {
 	// 0 %
 	subject := NewFixedSampler(0)
@@ -41,8 +23,17 @@ func TestFixedSampler(t *testing.T) {
 		}
 	}
 
-	rnd = newDeterministicGenerator()
 	// 50 %
+
+	// set seed for reproducibility.
+	rnd := rand.New(rand.NewSource(123))
+	intn = func(n int) int {
+		return rnd.Intn(n)
+	}
+	defer func() {
+		intn = rand.Intn
+	}()
+
 	trueCount := 0
 	subject = NewFixedSampler(33)
 	for i := 0; i < 100; i++ {
@@ -76,9 +67,17 @@ func TestAdaptiveSampler(t *testing.T) {
 		}
 	}
 
+	// set seed for reproducibility.
+	rnd := rand.New(rand.NewSource(123))
+	intn = func(n int) int {
+		return rnd.Intn(n)
+	}
+	defer func() {
+		intn = rand.Intn
+	}()
+
 	// change start time to 1s ago for a more predictable result.
 	trueCount := 0
-	rnd = newDeterministicGenerator()
 	now := time.Now()
 	subject.(*adaptiveSampler).start = now.Add(-time.Second)
 	for i := 99; i < 199; i++ {