forked from zeromicro/go-zero
/
tokenlimit.go
66 lines (58 loc) · 1.34 KB
/
tokenlimit.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
60
61
62
63
64
65
66
package main
import (
"flag"
"fmt"
"runtime"
"sync"
"sync/atomic"
"time"
"github.com/dylanNew/go-zero/core/limit"
"github.com/dylanNew/go-zero/core/stores/redis"
)
const (
burst = 100
rate = 100
seconds = 5
)
var (
rdx = flag.String("redis", "localhost:6379", "the redis, default localhost:6379")
rdxType = flag.String("redisType", "node", "the redis type, default node")
rdxKey = flag.String("redisKey", "rate", "the redis key, default rate")
rdxPass = flag.String("redisPass", "", "the redis password")
threads = flag.Int("threads", runtime.NumCPU(), "the concurrent threads, default to cores")
)
func main() {
flag.Parse()
store := redis.NewRedis(*rdx, *rdxType, *rdxPass)
fmt.Println(store.Ping())
limit := limit.NewTokenLimiter(rate, burst, store, *rdxKey)
timer := time.NewTimer(time.Second * seconds)
quit := make(chan struct{})
defer timer.Stop()
go func() {
<-timer.C
close(quit)
}()
var allowed, denied int32
var wait sync.WaitGroup
for i := 0; i < *threads; i++ {
wait.Add(1)
go func() {
for {
select {
case <-quit:
wait.Done()
return
default:
if limit.Allow() {
atomic.AddInt32(&allowed, 1)
} else {
atomic.AddInt32(&denied, 1)
}
}
}
}()
}
wait.Wait()
fmt.Printf("allowed: %d, denied: %d, qps: %d\n", allowed, denied, (allowed+denied)/seconds)
}