Skip to content
This repository was archived by the owner on Jun 14, 2023. It is now read-only.

Commit 14afcd9

Browse files
authored
fix(sampler): fixed random sampler called IsSampled will panic when concurrency call (#140) (#142)
1 parent 343dbb7 commit 14afcd9

File tree

2 files changed

+94
-4
lines changed

2 files changed

+94
-4
lines changed

sampler.go

+40-3
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"fmt"
2121
"math/rand"
2222
"strconv"
23+
"sync"
2324
"time"
2425
)
2526

@@ -44,20 +45,56 @@ func (s *ConstSampler) IsSampled(operation string) bool {
4445
return s.decision
4546
}
4647

48+
// RandomSampler Use sync.Pool to implement concurrent-safe for randomizer.
4749
type RandomSampler struct {
4850
samplingRate float64
49-
rand *rand.Rand
5051
threshold int
52+
pool sync.Pool
5153
}
5254

5355
// IsSampled implements IsSampled() of Sampler.
5456
func (s *RandomSampler) IsSampled(operation string) bool {
55-
return s.threshold > s.rand.Intn(100)
57+
58+
return s.threshold > s.generateRandomNumber()
5659
}
5760

5861
func (s *RandomSampler) init() {
59-
s.rand = rand.New(rand.NewSource(time.Now().Unix()))
62+
6063
s.threshold = int(s.samplingRate * 100)
64+
s.pool.New = s.newRand
65+
}
66+
67+
func (s *RandomSampler) generateRandomNumber() int {
68+
69+
r := s.getRandomizer()
70+
defer s.returnRandomizer(r)
71+
72+
return r.Intn(100)
73+
}
74+
75+
func (s *RandomSampler) returnRandomizer(r *rand.Rand) {
76+
s.pool.Put(r)
77+
}
78+
79+
func (s *RandomSampler) getRandomizer() *rand.Rand {
80+
81+
var r *rand.Rand
82+
83+
generator := s.pool.Get()
84+
if generator == nil {
85+
generator = s.newRand()
86+
}
87+
88+
r, ok := generator.(*rand.Rand)
89+
if !ok {
90+
r = s.newRand().(*rand.Rand) // it must be *rand.Rand
91+
}
92+
93+
return r
94+
}
95+
96+
func (s *RandomSampler) newRand() interface{} {
97+
return rand.New(rand.NewSource(time.Now().UnixNano()))
6198
}
6299

63100
func NewRandomSampler(samplingRate float64) *RandomSampler {

sampler_test.go

+54-1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package go2sky
1818

1919
import (
20+
"sync"
2021
"testing"
2122
)
2223

@@ -36,9 +37,16 @@ func TestConstSampler_IsSampled(t *testing.T) {
3637

3738
func TestRandomSampler_IsSampled(t *testing.T) {
3839
randomSampler := NewRandomSampler(0.5)
40+
41+
t.Run("threshold need transform", func(t *testing.T) {
42+
if randomSampler.threshold != 50 {
43+
t.Errorf("threshold should be 50")
44+
}
45+
})
46+
3947
operationName := "op"
4048

41-
//just for test case
49+
// just for test case
4250
randomSampler.threshold = 100
4351
sampled := randomSampler.IsSampled(operationName)
4452
if sampled != true {
@@ -51,3 +59,48 @@ func TestRandomSampler_IsSampled(t *testing.T) {
5159
t.Errorf("random sampler should not be sampled")
5260
}
5361
}
62+
63+
func TestNewRandomSampler(t *testing.T) {
64+
randomSampler := NewRandomSampler(100)
65+
operationName := "op"
66+
sampled := randomSampler.IsSampled(operationName)
67+
if sampled != true {
68+
t.Errorf("random sampler should be sampled")
69+
}
70+
}
71+
72+
func TestRandomSampler_getRandomizer(t *testing.T) {
73+
74+
t.Run("must not nil", func(t *testing.T) {
75+
76+
sampler := &RandomSampler{
77+
pool: sync.Pool{},
78+
}
79+
80+
if sampler.getRandomizer() == nil {
81+
t.Errorf("randomizer should be nil")
82+
}
83+
})
84+
85+
t.Run("must not nil, if got not a *rand.Rand", func(t *testing.T) {
86+
87+
sampler := &RandomSampler{
88+
pool: sync.Pool{},
89+
}
90+
91+
sampler.pool.Put(&struct{}{})
92+
if sampler.getRandomizer() == nil {
93+
t.Errorf("randomizer should be nil")
94+
}
95+
})
96+
}
97+
98+
func BenchmarkRandomPoolSampler_IsSampled(b *testing.B) {
99+
sampler := NewRandomSampler(0.5)
100+
operationName := "op"
101+
b.RunParallel(func(pb *testing.PB) {
102+
for pb.Next() {
103+
sampler.IsSampled(operationName)
104+
}
105+
})
106+
}

0 commit comments

Comments
 (0)