Go语言中进行API限流的实战详解

package main

import (
	"fmt"
	"time"
)

type TokenBucket struct {
	capacity int           // 桶的最大容量
	tokens   int           // 当前桶中令牌数量
	rate     time.Duration // 令牌生成速率
	lastTime time.Time     // 上次令牌生成时间
}

func NewTokenBucket(capacity int, rate time.Duration) *TokenBucket {
	return &TokenBucket{
		capacity: capacity,
		tokens:   capacity,
		rate:     rate,
		lastTime: time.Now(),
	}
}

func (tb *TokenBucket) Allow() bool {
	// 计算令牌的生成数量
	now := time.Now()
	elapsed := now.Sub(tb.lastTime)
	tokensToAdd := int(elapsed / tb.rate)
	tb.lastTime = now

	// 如果有令牌生成，更新桶中的令牌数量
	if tokensToAdd > 0 {
		tb.tokens = min(tb.capacity, tb.tokens+tokensToAdd)
	}

	// 如果有令牌，则消费一个令牌并返回 true
	if tb.tokens > 0 {
		tb.tokens--
		return true
	}
	// 否则返回 false，拒绝请求
	return false
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func main() {
	tb := NewTokenBucket(5, time.Second)

	// 模拟请求
	for i := 0; i < 10; i++ {
		if tb.Allow() {
			fmt.Println("Request", i, "allowed")
		} else {
			fmt.Println("Request", i, "denied")
		}
		time.Sleep(200 * time.Millisecond)
	}
}

package main

import (
	"fmt"
	"log"
	"time"

	"github.com/go-redis/redis/v8"
	"golang.org/x/net/context"
)

var rdb *redis.Client
var ctx = context.Background()

func init() {
	// 初始化 Redis 客户端
	rdb = redis.NewClient(&redis.Options{
		Addr: "localhost:6379", // Redis 地址
	})
}

func limitRequest(userID string) bool {
	// 使用 Redis 来存储访问记录
	key := fmt.Sprintf("rate_limit:%s", userID)
	// 设置每个用户的最大请求次数限制
	maxRequests := 5
	// 设置过期时间为 60 秒
	expiration := 60 * time.Second

	// 获取当前用户的请求次数
	count, err := rdb.Get(ctx, key).Int()
	if err != nil && err != redis.Nil {
		log.Fatalf("Error retrieving count from Redis: %v", err)
	}

	// 如果请求次数超过限制，拒绝请求
	if count >= maxRequests {
		return false
	}

	// 增加请求次数
	err = rdb.Incr(ctx, key).Err()
	if err != nil {
		log.Fatalf("Error incrementing count: %v", err)
	}

	// 设置过期时间，避免无限制的请求
	rdb.Expire(ctx, key, expiration)

	return true
}

func main() {
	userID := "user123"
	for i := 0; i < 10; i++ {
		if limitRequest(userID) {
			fmt.Println("Request allowed")
		} else {
			fmt.Println("Request denied due to rate limit")
		}
		time.Sleep(5 * time.Second)
	}
}