forked from fagongzi/manba
-
Notifications
You must be signed in to change notification settings - Fork 0
/
barrier.go
59 lines (49 loc) · 1.04 KB
/
barrier.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
package util
import (
"math/rand"
"sync"
"sync/atomic"
)
var (
lock sync.Mutex
randSources = make(map[int][]int)
)
func createRandSourceByBase(base int) []int {
lock.Lock()
defer lock.Unlock()
if value, ok := randSources[base]; ok {
return value
}
value := make([]int, base, base)
for i := 0; i < base; i++ {
value[i] = i
}
rand.Shuffle(base, func(i, j int) {
value[i], value[j] = value[j], value[i]
})
randSources[base] = value
return value
}
// RateBarrier rand barrier
type RateBarrier struct {
source []int
op uint64
rate int
base int
}
// NewRateBarrier returns a barrier based by 100
func NewRateBarrier(rate int) *RateBarrier {
return NewRateBarrierBase(rate, 100)
}
// NewRateBarrierBase returns a barrier with base
func NewRateBarrierBase(rate, base int) *RateBarrier {
return &RateBarrier{
source: createRandSourceByBase(base),
rate: rate,
base: base,
}
}
// Allow returns true if allowed
func (b *RateBarrier) Allow() bool {
return b.source[int(atomic.AddUint64(&b.op, 1))%b.base] < b.rate
}