├── go.mod ├── go.sum ├── .gitignore ├── benchark_test.go ├── ratelimit_test.go ├── readme.md └── ratelimit.go /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/maczam/ymratelimit 2 | 3 | go 1.14 4 | 5 | require github.com/juju/ratelimit v1.0.1 6 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY= 2 | github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk= 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | /bin/* 17 | /pkg/* 18 | 19 | 20 | 21 | # File-based project format 22 | *.ipr 23 | *.iml 24 | *.iws 25 | # Directory-based project format 26 | .idea/ -------------------------------------------------------------------------------- /benchark_test.go: -------------------------------------------------------------------------------- 1 | package ymratelimit 2 | 3 | import ( 4 | "github.com/juju/ratelimit" 5 | "testing" 6 | "time" 7 | ) 8 | 9 | func BenchmarkYmretelimit(b *testing.B) { 10 | rl := NewTokenBucket(time.Second, 15) // per second 11 | 12 | b.ResetTimer() 13 | 14 | for i := 0; i < b.N; i++ { 15 | rl.TakeAvailable() 16 | } 17 | } 18 | 19 | func BenchmarkParallelYmretelimit(b *testing.B) { 20 | rl := NewTokenBucket(time.Second, 15) // per second 21 | 22 | b.ResetTimer() 23 | 24 | b.RunParallel(func(pb *testing.PB) { 25 | 26 | for pb.Next() { 27 | rl.TakeAvailable() 28 | } 29 | }) 30 | } 31 | 32 | func BenchmarkJujuRatelimit(b *testing.B) { 33 | rl := ratelimit.NewBucket(time.Second, 15) 34 | b.ResetTimer() 35 | 36 | for i := 0; i < b.N; i++ { 37 | rl.TakeAvailable(1) 38 | } 39 | } 40 | 41 | func BenchmarkParallelJujuRatelimit(b *testing.B) { 42 | rl := ratelimit.NewBucket(time.Second, 15) 43 | b.ResetTimer() 44 | 45 | b.RunParallel(func(pb *testing.PB) { 46 | 47 | for pb.Next() { 48 | rl.TakeAvailable(1) 49 | } 50 | }) 51 | } 52 | -------------------------------------------------------------------------------- /ratelimit_test.go: -------------------------------------------------------------------------------- 1 | package ymratelimit 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | "sync/atomic" 7 | "testing" 8 | "time" 9 | ) 10 | 11 | func TestNewTokenBucket(t *testing.T) { 12 | 13 | rl := NewTokenBucket(time.Minute, 50) // per second 14 | 15 | startTime := time.Now() 16 | 17 | wg := sync.WaitGroup{} 18 | 19 | count := int32(0) 20 | stop := false 21 | for i := 0; i < 2; i++ { 22 | wg.Add(1) 23 | go func() { 24 | for !stop { 25 | if rl.TakeAvailable() { 26 | atomic.AddInt32(&count, 1) 27 | } 28 | time.Sleep(time.Millisecond * 100) 29 | } 30 | wg.Done() 31 | }() 32 | } 33 | 34 | time.Sleep(time.Second * 70) 35 | stop = true 36 | 37 | endTime := time.Now() 38 | 39 | wg.Wait() 40 | 41 | fmt.Println("t2与t1相差：", endTime.Sub(startTime), ";count：", count) //t2与t1相差： 50s 42 | } 43 | 44 | func TestNewLeakyBucket(t *testing.T) { 45 | 46 | rl := NewLeakyBucket(time.Second, 15) // per second 47 | 48 | startTime := time.Now() 49 | 50 | wg := sync.WaitGroup{} 51 | 52 | count := int32(0) 53 | 54 | stop := false 55 | for i := 0; i < 20; i++ { 56 | wg.Add(1) 57 | go func() { 58 | for !stop { 59 | if rl.TakeAvailable() { 60 | atomic.AddInt32(&count, 1) 61 | } 62 | //time.Sleep(time.Millisecond * 100) 63 | } 64 | wg.Done() 65 | }() 66 | } 67 | 68 | time.Sleep(time.Second * 5) 69 | stop = true 70 | 71 | endTime := time.Now() 72 | 73 | wg.Wait() 74 | 75 | fmt.Println("t2与t1相差：", endTime.Sub(startTime), ";count：", count) //t2与t1相差： 50s 76 | } 77 | 78 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # ymratelimit 2 | 3 | 目前在流量控制方面常用到的两个算法分别是，漏桶(Leaky bucket)[https://en.wikipedia.org/wiki/Leaky_bucket]与令牌桶(Token bucket)[https://en.wikipedia.org/wiki/Token_bucket]算法。 4 | 这两个算法在实现稍微有一点点不一样，链接里面wikipedia有详细解释，只是漏桶在实现过程中，会产生一个队列，允许瞬时并发。 5 | 6 | 为了提高并发，在性能方面优化: 7 | 8 | > * 禁止使用锁 9 | > * 每次请求最少只要一次cas操作 10 | > * 所有计数都转化成int64的操作，尽量减少cpu额外计算浪费 11 | 12 | 测试对比对象： github.com/juju/ratelimit 13 | ``` 14 | ➜ ymratelimit git:(master) ✗ go test -bench=. -run=none 15 | goos: darwin 16 | goarch: amd64 17 | pkg: github.com/maczam/ymratelimit 18 | BenchmarkYmretelimit-4 14109680 79.9 ns/op 19 | BenchmarkParallelYmretelimit-4 44515245 28.5 ns/op 20 | BenchmarkJujuRatelimit-4 10214019 111 ns/op 21 | BenchmarkParallelJujuRatelimit-4 6336103 160 ns/op 22 | PASS 23 | ok github.com/maczam/ymratelimit 4.978s 24 | 25 | ➜ ymratelimit git:(master) ✗ go test -bench=. -benchmem -run=none 26 | goos: darwin 27 | goarch: amd64 28 | pkg: github.com/maczam/ymratelimit 29 | BenchmarkYmretelimit-4 14484910 80.0 ns/op 0 B/op 0 allocs/op 30 | BenchmarkParallelYmretelimit-4 42125070 27.6 ns/op 0 B/op 0 allocs/op 31 | BenchmarkJujuRatelimit-4 10546452 111 ns/op 0 B/op 0 allocs/op 32 | BenchmarkParallelJujuRatelimit-4 6592738 171 ns/op 0 B/op 0 allocs/op 33 | PASS 34 | ok github.com/maczam/ymratelimit 5.034s 35 | 36 | ``` 37 | 38 | 单线程串行，差不多，但是多线程并发是JujuRatelimit性能7倍。 39 | 40 | 41 | # 使用 42 | > go get github.com/maczam/ymretelimit 43 | 44 | ## LeakyBucket 45 | ``` go 46 | rl := ymretelimit.NewLeakyBucket(time.Second, 15) // per second 47 | rl.TakeAvailable() 48 | ``` 49 | 50 | ## TokenBucket 51 | ``` go 52 | rl := ymretelimit.NewTokenBucket(time.Microsecond, 15) // per Microsecond 53 | rl.TakeAvailable() 54 | ``` 55 | 56 | ## LeakyBucket.TakeAvailable 57 | ``` go 58 | for !taken { 59 | var newLast int64 = 0 60 | previousStatePointer := atomic.LoadPointer(&t.lastTokenTimestamp) 61 | lastTokenTimestamp := (*int64)(previousStatePointer) 62 | // 本次需要需要到达时间 当前时间戳-上次获取的时间戳 + 每次请求时间片段 63 | newLast = *lastTokenTimestamp + t.perRequest 64 | //sb.WriteString(fmt.Sprintf("lastTokenTimestamp %d;newLast:%v;", *lastTokenTimestamp, newLast)) 65 | 66 | if now < newLast { 67 | break 68 | } else { 69 | // 如果下一个线程 70 | taken = atomic.CompareAndSwapPointer(&t.lastTokenTimestamp, previousStatePointer, unsafe.Pointer(&newLast)) 71 | } 72 | } 73 | ``` 74 | ## TokenBucket.TakeAvailable 75 | 76 | 注解代码 77 | ``` go 78 | func (t *tokenBucket) TakeAvailableWithNow(now int64) bool { 79 | 80 | // for 是为了保证LoadPointer和CompareAndSwapPointer是处于原子状态 81 | taken := false 82 | 83 | for !taken { 84 | lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat) 85 | lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer) 86 | 87 | if now > lastTokenBucketStat.nextTokenTimestamp { 88 | newStat := tokenBucketStat{} 89 | newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp + t.fillInterval 90 | newStat.keepCapacity = t.capacity - 1 91 | taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat)) 92 | } else { 93 | 94 | // 已经没有了 95 | if lastTokenBucketStat.keepCapacity > 0 { 96 | newStat := tokenBucketStat{} 97 | newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp 98 | newStat.keepCapacity = lastTokenBucketStat.keepCapacity - 1 99 | taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat)) 100 | } else { 101 | break 102 | } 103 | } 104 | } 105 | return taken 106 | } 107 | ``` -------------------------------------------------------------------------------- /ratelimit.go: -------------------------------------------------------------------------------- 1 | package ymratelimit 2 | 3 | import ( 4 | "encoding/json" 5 | "sync/atomic" 6 | "time" 7 | "unsafe" 8 | ) 9 | 10 | /** 11 | 抽象接口 12 | */ 13 | type Limiter interface { 14 | TakeAvailableWithNow(now int64) bool 15 | TakeAvailable() bool 16 | GetCapacity() int64 17 | GetLegacyCapacity() int64 18 | } 19 | 20 | /** 21 | 漏桶算法能够强行限制数据的传输速率。 22 | */ 23 | type leakyBucket struct { 24 | capacity int64 25 | fillInterval int64 //统计周期 26 | lastTokenTimestamp unsafe.Pointer 27 | perRequest int64 //计算出每次token占用的时间片段 28 | } 29 | 30 | /** 31 | 如果能获取，那么不用判断时间， time.Now().UnixNano() 必须使用UnixNano 32 | */ 33 | func (t *leakyBucket) TakeAvailableWithNow(now int64) bool { 34 | 35 | // for 是为了保证LoadPointer和CompareAndSwapPointer是处于原子状态 36 | taken := false 37 | //sb := strings.Builder{} 38 | //t2 := time.Unix(0, now) 39 | 40 | //sb.WriteString(fmt.Sprintf("now %d;nowTime:%v;PerRequest:%d;", now, t2, t.PerRequest)) 41 | for !taken { 42 | var newLast int64 = 0 43 | previousStatePointer := atomic.LoadPointer(&t.lastTokenTimestamp) 44 | lastTokenTimestamp := (*int64)(previousStatePointer) 45 | // 本次需要需要到达时间 当前时间戳-上次获取的时间戳 + 每次请求时间片段 46 | newLast = *lastTokenTimestamp + t.perRequest 47 | //sb.WriteString(fmt.Sprintf("lastTokenTimestamp %d;newLast:%v;", *lastTokenTimestamp, newLast)) 48 | 49 | if now < newLast { 50 | break 51 | //sb.WriteString("now < newLast;") 52 | } else { 53 | // 如果下一个线程 54 | //sb.WriteString("now < newLast;") 55 | taken = atomic.CompareAndSwapPointer(&t.lastTokenTimestamp, previousStatePointer, unsafe.Pointer(&newLast)) 56 | } 57 | } 58 | //sb.WriteString(fmt.Sprintf("最终结果:%t", taken)) 59 | //fmt.Println(sb.String()) 60 | return taken 61 | } 62 | 63 | func (t *leakyBucket) TakeAvailable() bool { 64 | return t.TakeAvailableWithNow(time.Now().UnixNano()) 65 | } 66 | 67 | func (t *leakyBucket) GetCapacity() int64 { 68 | return t.capacity 69 | } 70 | func (t *leakyBucket) GetLegacyCapacity() int64 { 71 | return -1 72 | } 73 | 74 | func (t *leakyBucket) MarshalJSON() ([]byte, error) { 75 | object := map[string]interface{}{} 76 | object["capacity"] = t.capacity 77 | return json.Marshal(object) 78 | } 79 | 80 | func NewLeakyBucket(fillInterval time.Duration, capacity int64) Limiter { 81 | fillIntervalInt := int64(fillInterval) 82 | l := &leakyBucket{ 83 | fillInterval: fillIntervalInt, 84 | perRequest: fillIntervalInt / capacity, 85 | capacity: capacity, 86 | } 87 | lastTokenTimestamp := time.Now().UnixNano() 88 | l.lastTokenTimestamp = unsafe.Pointer(&lastTokenTimestamp) 89 | return l 90 | } 91 | 92 | /** 93 | 令牌桶算法能够在限制数据的平均传输速率的同时还允许某种程度的突发传输。 94 | */ 95 | type tokenBucket struct { 96 | capacity int64 97 | fillInterval int64 //统计周期 98 | tokenBucketStat unsafe.Pointer //当前状态 99 | perRequest int64 //计算出每次token占用的时间片段 100 | } 101 | 102 | /** 103 | 当前状态 104 | */ 105 | type tokenBucketStat struct { 106 | nextTokenTimestamp int64 107 | keepCapacity int64 //本次time window 还剩下多少次 108 | } 109 | 110 | /** 111 | 如果能获取，那么不用判断时间， time.Now().UnixNano() 必须使用UnixNano 112 | */ 113 | func (t *tokenBucket) TakeAvailableWithNow(now int64) bool { 114 | 115 | // for 是为了保证LoadPointer和CompareAndSwapPointer是处于原子状态 116 | taken := false 117 | //sb := strings.Builder{} 118 | //t2 := time.Unix(0, now) 119 | // 120 | //sb.WriteString(fmt.Sprintf("now %d;nowTime:%v", now, t2)) 121 | for !taken { 122 | lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat) 123 | lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer) 124 | 125 | //sb.WriteString(fmt.Sprintf("下个窗口时间:%d;", lastTokenBucketStat.NextTokenTimestamp)) 126 | //sb.WriteString(fmt.Sprintf("距离下一次:%d;", now-lastTokenBucketStat.NextTokenTimestamp)) 127 | //sb.WriteString(fmt.Sprintf("lastKeepCapacity:%d;", lastTokenBucketStat.KeepCapacity)) 128 | 129 | if now > lastTokenBucketStat.nextTokenTimestamp { 130 | newStat := tokenBucketStat{} 131 | newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp + t.fillInterval 132 | newStat.keepCapacity = t.capacity - 1 133 | //sb.WriteString(fmt.Sprintf("改写下一次时间:%d;下一次容量:%d;", newStat.NextTokenTimestamp, newStat.KeepCapacity)) 134 | taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat)) 135 | } else { 136 | //sb.WriteString(fmt.Sprintf("在时间窗口之内;")) 137 | 138 | // 已经没有了 139 | if lastTokenBucketStat.keepCapacity > 0 { 140 | newStat := tokenBucketStat{} 141 | newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp 142 | newStat.keepCapacity = lastTokenBucketStat.keepCapacity - 1 143 | //sb.WriteString(fmt.Sprintf(fmt.Sprintf("修改结余:%d;", newStat.KeepCapacity))) 144 | taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat)) 145 | } else { 146 | break 147 | } 148 | } 149 | } 150 | //sb.WriteString(fmt.Sprintf("最终结果:%t", taken)) 151 | //fmt.Println(sb.String()) 152 | return taken 153 | } 154 | 155 | func (t *tokenBucket) TakeAvailable() bool { 156 | return t.TakeAvailableWithNow(time.Now().UnixNano()) 157 | } 158 | 159 | func (t *tokenBucket) GetCapacity() int64 { 160 | return t.capacity 161 | } 162 | 163 | func (t *tokenBucket) GetLegacyCapacity() int64 { 164 | lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat) 165 | lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer) 166 | return lastTokenBucketStat.keepCapacity 167 | } 168 | 169 | func (t *tokenBucket) MarshalJSON() ([]byte, error) { 170 | object := map[string]interface{}{} 171 | object["capacity"] = t.capacity 172 | lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat) 173 | lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer) 174 | object["keepCapacity"] = lastTokenBucketStat.keepCapacity 175 | return json.Marshal(object) 176 | } 177 | 178 | /** 179 | 令牌桶算法能够在限制数据的平均传输速率的同时还允许某种程度的突发传输。 180 | */ 181 | func NewTokenBucket(fillInterval time.Duration, capacity int64) Limiter { 182 | fillIntervalInt := int64(fillInterval) 183 | l := &tokenBucket{ 184 | fillInterval: fillIntervalInt, 185 | capacity: capacity, 186 | } 187 | tokenBucketStat := tokenBucketStat{ 188 | nextTokenTimestamp: time.Now().UnixNano(), 189 | keepCapacity: capacity, 190 | } 191 | l.tokenBucketStat = unsafe.Pointer(&tokenBucketStat) 192 | return l 193 | } 194 | --------------------------------------------------------------------------------