go - 令牌桶

最近在看 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)
}