Skip to content

Commit

Permalink
Fix/trace sampler (#3467)
Browse files Browse the repository at this point in the history
* Revert "Use crypto/rand for sampling (#3466)"

This reverts commit a17e687.

* Use global random generator which is concurrency-safe
  • Loading branch information
raphael authored Jan 27, 2024
1 parent a17e687 commit 3bbb4b2
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 53 deletions.
44 changes: 11 additions & 33 deletions middleware/sampler.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
package middleware

import (
"crypto/rand"
"math/big"
"math/rand"
"sync"
"sync/atomic"
"time"
Expand All @@ -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
Expand All @@ -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.
//
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
}
39 changes: 19 additions & 20 deletions middleware/sampler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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++ {
Expand Down Expand Up @@ -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++ {
Expand Down

0 comments on commit 3bbb4b2

Please sign in to comment.