From 57a614077fded7bf536545c109266b868be1d39f Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Tue, 16 Jul 2019 15:29:18 +0800
Subject: [PATCH] avoid data race while using Rand in sre_breaker.go

---
 pkg/net/netutil/breaker/sre_breaker.go      | 26 +++++++++++++++------
 pkg/net/netutil/breaker/sre_breaker_test.go | 18 ++++++++++++--
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/pkg/net/netutil/breaker/sre_breaker.go b/pkg/net/netutil/breaker/sre_breaker.go
index 5e410cd98..ef4c32856 100644
--- a/pkg/net/netutil/breaker/sre_breaker.go
+++ b/pkg/net/netutil/breaker/sre_breaker.go
@@ -3,6 +3,7 @@ package breaker
 import (
 	"math"
 	"math/rand"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -11,6 +12,12 @@ import (
 	"github.com/bilibili/kratos/pkg/stat/metric"
 )
 
+var (
+	// rand.New(...) returns a non thread safe object
+	random = rand.New(rand.NewSource(time.Now().UnixNano()))
+	lock   sync.Mutex
+)
+
 // sreBreaker is a sre CircuitBreaker pattern.
 type sreBreaker struct {
 	stat metric.RollingCounter
@@ -19,7 +26,6 @@ type sreBreaker struct {
 	request int64
 
 	state int32
-	r     *rand.Rand
 }
 
 func newSRE(c *Config) Breaker {
@@ -30,7 +36,6 @@ func newSRE(c *Config) Breaker {
 	stat := metric.NewRollingCounter(counterOpts)
 	return &sreBreaker{
 		stat: stat,
-		r:    rand.New(rand.NewSource(time.Now().UnixNano())),
 
 		request: c.Request,
 		k:       c.K,
@@ -69,14 +74,14 @@ func (b *sreBreaker) Allow() error {
 		atomic.CompareAndSwapInt32(&b.state, StateClosed, StateOpen)
 	}
 	dr := math.Max(0, (float64(total)-k)/float64(total+1))
-	rr := b.r.Float64()
+	drop := trueOnProba(dr)
 	if log.V(5) {
-		log.Info("breaker: drop ratio: %f, real rand: %f, drop: %v", dr, rr, dr > rr)
+		log.Info("breaker: drop ratio: %f, drop: %t", dr, drop)
 	}
-	if dr <= rr {
-		return nil
+	if drop {
+		return ecode.ServiceUnavailable
 	}
-	return ecode.ServiceUnavailable
+	return nil
 }
 
 func (b *sreBreaker) MarkSuccess() {
@@ -88,3 +93,10 @@ func (b *sreBreaker) MarkFailed() {
 	// drop ratio higher.
 	b.stat.Add(0)
 }
+
+func trueOnProba(proba float64) (truth bool) {
+	lock.Lock()
+	truth = random.Float64() < proba
+	lock.Unlock()
+	return
+}
diff --git a/pkg/net/netutil/breaker/sre_breaker_test.go b/pkg/net/netutil/breaker/sre_breaker_test.go
index 4559c4fc5..bd1849d74 100644
--- a/pkg/net/netutil/breaker/sre_breaker_test.go
+++ b/pkg/net/netutil/breaker/sre_breaker_test.go
@@ -5,7 +5,7 @@ import (
 	"math/rand"
 	"testing"
 	"time"
-	
+
 	"github.com/bilibili/kratos/pkg/stat/metric"
 	xtime "github.com/bilibili/kratos/pkg/time"
 
@@ -29,7 +29,6 @@ func getSREBreaker() *sreBreaker {
 	stat := metric.NewRollingCounter(counterOpts)
 	return &sreBreaker{
 		stat: stat,
-		r:    rand.New(rand.NewSource(time.Now().UnixNano())),
 
 		request: 100,
 		k:       2,
@@ -147,6 +146,21 @@ func TestSRESummary(t *testing.T) {
 	})
 }
 
+func TestTrueOnProba(t *testing.T) {
+	const proba = math.Pi / 10
+	const total = 100000
+	const epsilon = 0.05
+	var count int
+	for i := 0; i < total; i++ {
+		if trueOnProba(proba) {
+			count++
+		}
+	}
+
+	ratio := float64(count) / float64(total)
+	assert.InEpsilon(t, proba, ratio, epsilon)
+}
+
 func BenchmarkSreBreakerAllow(b *testing.B) {
 	breaker := getSRE()
 	b.ResetTimer()