一、令牌桶算法(控制速率)
Redis可以通过实现一个基于令牌桶算法的限流器来限制HTTP请求的访问速率。令牌桶算法是一种常见的限流算法,它基于一个令牌桶来控制请求的速率。
具体地,我们可以使用Redis中的有序集合(sorted set)来实现令牌桶算法。首先,我们需要在Redis中设置一个有序集合,将时间戳作为成员(member),将令牌数作为分值(score)。然后,每次收到一个请求时,就从有序集合中获取当前的令牌数,并将其减1。如果令牌数为0,表示该请求已超出限速,应该拒绝该请求。同时,我们也需要根据一定的规则,定期向有序集合中添加令牌,确保令牌桶中的令牌数不会永久耗尽。
以下是示例代码:
package main
import (
"context"
"fmt"
"time"
"github.com/go-redis/redis/v8"
)
func main() {
// Redis客户端和上下文
client := redis.NewClient(&redis.Options{
Addr: "localhost:6379",
Password: "",
DB: 0,
})
ctx := context.Background()
// 限速器的前缀
prefix := "rate_limit:"
// 每秒钟最多处理的请求数
rate := 5
// 限速器的时间间隔
interval := 1 * time.Second
// 处理请求的函数
processRequest := func(i int) {
fmt.Printf("Processing request #%d\n", i)
}
// 限速器函数
limit := func(key string) bool {
now := time.Now().UnixNano()
// 移除时间戳小于(now - interval)的所有成员
client.ZRemRangeByScore(ctx, key, "0", fmt.Sprintf("%d", now-int64(interval)))
// 获取当前集合长度,即令牌数
size, _ := client.ZCard(ctx, key).Result()
if size < int64(rate) {
// 如果令牌数小于最大请求数,则添加一个新的时间戳成为成员
client.ZAdd(ctx, key, &redis.Z{Score: float64(now), Member: float64(now)})
return true
}
// 否则拒绝该请求
fmt.Println("Rate limited")
return false
}
// lua脚本版限速器函数,可以保证原子性操作
// limit := func(key string) bool {
//
// // 定义lua脚本
// script := `
// local key = KEYS[1]
// local min = ARGV[1]
// local max = ARGV[2]
// local now = ARGV[3]
// local rate = tonumber(ARGV[4])
// redis.call("ZREMRANGEBYSCORE", key, min, max)
// local size = tonumber(redis.call("ZCARD", key))
// if size < rate then
// redis.call("ZADD", key, now, now)
// return 1
// end
// return 0
//`
// //当前时间 单位:纳秒
// now := time.Now().UnixNano()
//
// // 执行lua脚本
// result, err := client.Eval(ctx, script, []string{key}, "0", now-int64(interval), now, rate).Result()
// if err != nil {
// fmt.Println(err)
// return false
// } else if result == int64(1) {
// return true
// } else {
// fmt.Println("Rate limited...")
// return false
// }
// }
// 模拟处理100个请求
for i := 0; i < 100; i++ {
key := prefix + "test"
if ok := limit(key); ok { // 调用限速器
processRequest(i)
}
time.Sleep(100 * time.Millisecond) // 模拟请求间隔
}
}
执行效果:
go run main.go
redis储存token情况如下:
最多5个令牌数量,因为每次请求打过来都会根据速率参数将过期的token删除。
二、固定时间固定请求次数:
package main
import (
"fmt"
"time"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
)
var rdb *redis.Client
func init() {
rdb = redis.NewClient(&redis.Options{
Addr: "localhost:6379",
Password: "",
DB: 0,
})
}
func main() {
r := gin.Default()
r.Use(RateLimiter(time.Second*15, 1))
r.GET("/", func(c *gin.Context) {
c.JSON(200, gin.H{
"message": "Hello World!",
})
})
r.Run(":8080")
}
func RateLimiter(duration time.Duration, limit int64) gin.HandlerFunc {
return func(c *gin.Context) {
ip := c.ClientIP()
key := fmt.Sprintf("%s:%d", ip, time.Now().Unix()/int64(duration.Seconds()))
// 使用 Lua 脚本原子性操作
script := `
local current
current = redis.call("INCR", KEYS[1])
if current == 1 then
redis.call("EXPIRE", KEYS[1], ARGV[1])
end
return current
`
count, err := rdb.Eval(c.Request.Context(), script, []string{key}, int64(duration.Seconds())).Int64()
if err != nil {
c.AbortWithStatusJSON(500, gin.H{"error": err.Error()})
return
}
if count > limit {
c.AbortWithStatusJSON(429, gin.H{"error": "Too Many Requests"})
return
}
c.Next()
}
}
以上代码演示了在15秒内只能有1次请求的限制,并将这个规则做成了gin框架的中间件。