fix(sampler): fixed random sampler called IsSampled will panic when concurrency call (#140) (#142)

Author: Luckyboys

sampler.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"math/rand"
 	"strconv"
+	"sync"
 	"time"
 )
 
@@ -44,20 +45,56 @@ func (s *ConstSampler) IsSampled(operation string) bool {
 	return s.decision
 }
 
+// RandomSampler Use sync.Pool to implement concurrent-safe for randomizer.
 type RandomSampler struct {
 	samplingRate float64
-	rand         *rand.Rand
 	threshold    int
+	pool         sync.Pool
 }
 
 // IsSampled implements IsSampled() of Sampler.
 func (s *RandomSampler) IsSampled(operation string) bool {
-	return s.threshold > s.rand.Intn(100)
+
+	return s.threshold > s.generateRandomNumber()
 }
 
 func (s *RandomSampler) init() {
-	s.rand = rand.New(rand.NewSource(time.Now().Unix()))
+
 	s.threshold = int(s.samplingRate * 100)
+	s.pool.New = s.newRand
+}
+
+func (s *RandomSampler) generateRandomNumber() int {
+
+	r := s.getRandomizer()
+	defer s.returnRandomizer(r)
+
+	return r.Intn(100)
+}
+
+func (s *RandomSampler) returnRandomizer(r *rand.Rand) {
+	s.pool.Put(r)
+}
+
+func (s *RandomSampler) getRandomizer() *rand.Rand {
+
+	var r *rand.Rand
+
+	generator := s.pool.Get()
+	if generator == nil {
+		generator = s.newRand()
+	}
+
+	r, ok := generator.(*rand.Rand)
+	if !ok {
+		r = s.newRand().(*rand.Rand) // it must be *rand.Rand
+	}
+
+	return r
+}
+
+func (s *RandomSampler) newRand() interface{} {
+	return rand.New(rand.NewSource(time.Now().UnixNano()))
 }
 
 func NewRandomSampler(samplingRate float64) *RandomSampler {

sampler_test.go

@@ -17,6 +17,7 @@
 package go2sky
 
 import (
+	"sync"
 	"testing"
 )
 
@@ -36,9 +37,16 @@ func TestConstSampler_IsSampled(t *testing.T) {
 
 func TestRandomSampler_IsSampled(t *testing.T) {
 	randomSampler := NewRandomSampler(0.5)
+
+	t.Run("threshold need transform", func(t *testing.T) {
+		if randomSampler.threshold != 50 {
+			t.Errorf("threshold should be 50")
+		}
+	})
+
 	operationName := "op"
 
-	//just for test case
+	// just for test case
 	randomSampler.threshold = 100
 	sampled := randomSampler.IsSampled(operationName)
 	if sampled != true {
@@ -51,3 +59,48 @@ func TestRandomSampler_IsSampled(t *testing.T) {
 		t.Errorf("random sampler should not be sampled")
 	}
 }
+
+func TestNewRandomSampler(t *testing.T) {
+	randomSampler := NewRandomSampler(100)
+	operationName := "op"
+	sampled := randomSampler.IsSampled(operationName)
+	if sampled != true {
+		t.Errorf("random sampler should be sampled")
+	}
+}
+
+func TestRandomSampler_getRandomizer(t *testing.T) {
+
+	t.Run("must not nil", func(t *testing.T) {
+
+		sampler := &RandomSampler{
+			pool: sync.Pool{},
+		}
+
+		if sampler.getRandomizer() == nil {
+			t.Errorf("randomizer should be nil")
+		}
+	})
+
+	t.Run("must not nil, if got not a *rand.Rand", func(t *testing.T) {
+
+		sampler := &RandomSampler{
+			pool: sync.Pool{},
+		}
+
+		sampler.pool.Put(&struct{}{})
+		if sampler.getRandomizer() == nil {
+			t.Errorf("randomizer should be nil")
+		}
+	})
+}
+
+func BenchmarkRandomPoolSampler_IsSampled(b *testing.B) {
+	sampler := NewRandomSampler(0.5)
+	operationName := "op"
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			sampler.IsSampled(operationName)
+		}
+	})
+}