最近在看 go-kit, 涉及到限流时看到这个, 记录一下实现
后期补充一个 http 服务的限流例子
安装 rate 包: go get golang.org/x/time/rate
package main
import (
"context"
"fmt"
"golang.org/x/time/rate"
"time"
)
// Wait/WaitN 尝试取出, 成功继续执行, 不成功阻塞, 直到取出成功
func testWait(x rate.Limit, y, z int) {
// y: 令牌桶总容量
// x: 每秒生成令牌个数
// z: 每次取出令牌个数
c := context.Background()
// 创建令牌桶
r := rate.NewLimiter(x, y)
for true {
// 尝试取出 z 个, 成功继续执行, 不成功阻塞
r.WaitN(c, z) // r.Wait(c) 等价 r.WaitN(c, 1)
fmt.Println(time.Now().Format("2006-01-02 15:04:05"))
time.Sleep(time.Second)
}
}
// Allow/AllowN 尝试取出, 成功执行 A, 不成功试行 B
func testAllow(x rate.Limit, y, z int) {
// y: 令牌桶总容量
// x: 每秒生成令牌个数
// z: 每次取出令牌个数
// 创建令牌桶
r := rate.NewLimiter(x, y)
for true {
// 尝试取出 z 个, 如果成功执行 A, 不成功执行 B
if r.AllowN(time.Now(), z) { // r.Allow() 等价 r.WaitN(time.Now(), 1)
fmt.Println("正常逻辑代码")
} else {
fmt.Println("限流逻辑代码")
}
time.Sleep(time.Second)
}
}
func main() {
// Wait 尝试取出, 成功继续执行, 不成功阻塞, 直到取出成功
testWait(1, 5, 5)
// Allow 尝试取出, 成功执行 A, 不成功试行 B
testAllow(1, 5, 5)
}