Skip to content

Commit

Permalink
Use crypto/rand for sampling (#3466)
Browse files Browse the repository at this point in the history
  • Loading branch information
nitinmohan87 authored Jan 26, 2024
1 parent 8142df9 commit a17e687
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 11 deletions.
41 changes: 32 additions & 9 deletions middleware/sampler.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
package middleware

import (
"math/rand"
"crypto/rand"
"math/big"
"sync"
"sync/atomic"
"time"
Expand All @@ -14,6 +15,16 @@ 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 @@ -28,11 +39,11 @@ type (

const (
// adaptive upper bound has granularity in case caller becomes extremely busy.
adaptiveUpperBoundInt = 10000
adaptiveUpperBoundInt = int64(10000)
adaptiveUpperBoundFloat = float64(adaptiveUpperBoundInt)
)

var rnd = rand.New(rand.NewSource(1))
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 @@ -68,14 +79,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 int
var currentRate int64
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 = int((float64(s.maxSamplingRate) * adaptiveUpperBoundFloat) / r)
currentRate = int64((float64(s.maxSamplingRate) * adaptiveUpperBoundFloat) / r)
if currentRate > adaptiveUpperBoundInt {
currentRate = adaptiveUpperBoundInt
} else if currentRate < 1 {
Expand All @@ -86,15 +97,27 @@ func (s *adaptiveSampler) Sample() bool {
s.Unlock()
atomic.StoreInt64(&s.lastRate, int64(currentRate))
} else {
currentRate = int(atomic.LoadInt64(&s.lastRate))
currentRate = int64(atomic.LoadInt64(&s.lastRate))
}

// currentRate is never zero.
return currentRate == adaptiveUpperBoundInt || rnd.Intn(adaptiveUpperBoundInt) < currentRate
return currentRate == adaptiveUpperBoundInt || rnd.Int64(adaptiveUpperBoundInt) < currentRate
}

// Sample implementation for fixed percentage
func (s fixedSampler) Sample() bool {
samplingPercent := int(s)
return samplingPercent > 0 && (samplingPercent == 100 || rnd.Intn(100) < samplingPercent)
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()
}
23 changes: 21 additions & 2 deletions middleware/sampler_test.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,29 @@
package middleware

import (
"math/rand"
"testing"
"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 @@ -22,8 +41,8 @@ func TestFixedSampler(t *testing.T) {
}
}

rnd = newDeterministicGenerator()
// 50 %
rnd.Seed(123) // set seed for reproducibility.
trueCount := 0
subject = NewFixedSampler(33)
for i := 0; i < 100; i++ {
Expand Down Expand Up @@ -59,7 +78,7 @@ func TestAdaptiveSampler(t *testing.T) {

// change start time to 1s ago for a more predictable result.
trueCount := 0
rnd.Seed(123) // set seed for reproducibility.
rnd = newDeterministicGenerator()
now := time.Now()
subject.(*adaptiveSampler).start = now.Add(-time.Second)
for i := 99; i < 199; i++ {
Expand Down

0 comments on commit a17e687

Please sign in to comment.