Golang Rate - 令牌桶限流包使用

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

// - r (rate.Limit): 这是令牌的生成速率，表示每秒允许的事件数
// - b (int): 这是令牌桶的容量（burst），表示桶中最多可以存储多少个令牌。它决定了限流器可以处理的突发请求数量。
func NewLimiter(r Limit, b int) *Limiter

// 非阻塞。如果桶中有足够的令牌，它会消耗一个令牌并返回 true；否则，返回 false。
func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool

// 阻塞。如果桶中没有令牌，它会阻塞当前 goroutine，等待直到有令牌可用。
// ctx 参数:
// - 如果 ctx 被取消（例如，context.WithCancel 或 context.WithTimeout），Wait 会立即返回 ctx.Err()。
// - 如果 ctx 为 nil，则不会有超时或取消，Wait 会一直等待直到令牌可用。
func (lim *Limiter) Wait(ctx context.Context) error
func (lim *Limiter) WaitN(ctx context.Context, n int) error

// 这些方法用于预留令牌，但不会立即消耗。
func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

// 动态调整速率
limiter.SetLimit(rate.Limit(20)) // 调整为每秒20个
limiter.SetBurst(10)             // 调整容量为10

 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

import (
	"fmt"
	"time"

	"golang.org/x/time/rate"
)

func main() {
	// 创建一个限流器：每秒允许10个事件，桶容量为5
	// 这意味着：
	// - 长期平均速率不会超过每秒10个请求。
	// - 短时间内，可以处理最多5个突发请求（如果桶是满的）。
	limiter := rate.NewLimiter(rate.Limit(10), 5)

	fmt.Printf("Limiter created: Rate=%v, Burst=%v\n", limiter.Limit(), limiter.Burst())

	// 尝试获取令牌
	// Allow() 是非阻塞的，消耗一个token，立即返回是否允许
	if limiter.Allow() {
		fmt.Println("Request 1 allowed immediately.")
	} else {
		fmt.Println("Request 1 denied immediately.")
	}

	// Wait() 是阻塞的，直到获取到令牌
	fmt.Println("Waiting for a token...")
	err := limiter.Wait(nil) // nil context表示没有超时或取消
	if err != nil {
		fmt.Printf("Wait error: %v\n", err)
	} else {
		fmt.Println("Request 2 allowed after waiting (if necessary).")
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

// 限制对某个第三方API的调用，例如每秒100次
apiLimiter := rate.NewLimiter(rate.Limit(100), 100)

func callExternalAPI() {
    if err := apiLimiter.Wait(context.Background()); err != nil {
        fmt.Println("Failed to acquire token for API call:", err)
        return
    }
    // 实际调用API的逻辑
    fmt.Println("Calling external API...")
}

// 在goroutine中并发调用
for i := 0; i < 200; i++ {
    go callExternalAPI()
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

// 限制服务器处理请求的速率，例如每秒500个请求
requestLimiter := rate.NewLimiter(rate.Limit(500), 500)

func handleRequest(req interface{}) {
    if !requestLimiter.Allow() {
        // 请求过多，直接拒绝或返回错误
        fmt.Println("Request denied: Too many requests.")
        return
    }
    // 处理请求的逻辑
    fmt.Println("Processing request...")
}

// 在HTTP服务器的Handler中
// http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
//     if !requestLimiter.Allow() {
//         http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
//         return
//     }
//     // ... 处理请求 ...
// })

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

// 限制后台任务的执行速率，例如每分钟10个任务
taskLimiter := rate.NewLimiter(rate.Every(time.Minute/10), 1)

func processTask(taskID int) {
    fmt.Printf("Processing task %d...\n", taskID)
    time.Sleep(50 * time.Millisecond) // 模拟任务耗时
}

func main() {
    fmt.Println("\n--- Background Task Limiting ---")
    for i := 0; i < 20; i++ {
        // 阻塞等待令牌，确保任务不会过快执行
        if err := taskLimiter.Wait(context.Background()); err != nil {
            fmt.Printf("Error waiting for task %d: %v\n", i, err)
            continue
        }
        go processTask(i)
    }
    time.Sleep(2 * time.Minute) // 等待所有任务完成
}