大家好，又见面了，我是你们的朋友全栈君。
主要思路：
1.创建channel，通过定时器定时往channel中写入令牌，返回令牌桶本身（channel）；
2.判断请求是否可以拿到令牌；
package main
import (
"context"
"fmt"
"time"
)
func getBucket(capacityPs, maxCapacity int) (chan struct{}, *time.Ticker) {
var bucketToken = make(chan struct{}, maxCapacity)
ticker := time.NewTicker(time.Second / time.Duration(capacityPs))
go func() {
for {
select {
case <-ticker.C:
bucketToken <- struct{}{}
default:
}
}
}()
return bucketToken, ticker
}
func getToken(ctx context.Context, bucket *chan struct{}, wait bool) bool {
if wait {
select {
case <-*bucket:
return true
case <-ctx.Done():
fmt.Println("请求超时...")
return false
}
} else { // 不阻塞
select {
case <-*bucket:
return true
default:
return false
}
}
}
// 模拟qps 100， 超时等待时间为1s，模拟1000个请求进来
// 令牌桶最大每秒产生30个令牌，最大突发为50（桶最大容量）
func test1() {
bucket, ticker := getBucket(30, 50)
i := 0
total := 0
fail := 0
go func() {
for { // qps 100
i++
time.Sleep(time.Millisecond * 10)
ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*100)
ok := getToken(ctx, &bucket, false)
if ok {
total++
} else {
fail++
}
fmt.Println("当前获取token是否成功：", ok)
if i >= 1000 {
ticker.Stop()
break
}
}
}()
time.Sleep(time.Second * 20)
fmt.Println(i, total, fail)
}
// 模拟qps 100， 超时等待时间为1s，模拟1000个请求进来
// 令牌桶最大每秒产生30个令牌，最大突发为50（桶最大容量）
func test2() {
bucket, ticker := getBucket(30, 50)
i := 0
total := 0
fail := 0
go func() {
for { // qps 100
i++
time.Sleep(time.Millisecond * 10)
ctx, _ := context.WithTimeout(context.Background(), time.Millisecond*20)
ok := getToken(ctx, &bucket, true)
if ok {
total++
} else {
fail++
}
fmt.Println("当前获取token是否成功：", ok)
if i >= 1000 {
ticker.Stop()
break
}
}
}()
time.Sleep(time.Second * 20)
fmt.Println(i, total, fail)
}
func main() {
//test1()
test2()
}
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容， 请联系我们举报，一经查实，本站将立刻删除。
发布者：全栈程序员-站长，转载请注明出处：https://javaforall.net/155395.html原文链接：https://javaforall.net
