redis实现基于令牌桶算法(控制速率)以及固定时间的固定请求次数的请求限流器(golang版)

一、令牌桶算法(控制速率)

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 int) gin.HandlerFunc {
    return func(c *gin.Context) {
        ip := c.ClientIP()
        key := fmt.Sprintf("%s:%d", ip, time.Now().Unix()/int64(duration.Seconds()))

        count, err := rdb.Incr(c.Request.Context(), key).Result()
        if err != nil {
            c.AbortWithStatusJSON(500, gin.H{"error": err.Error()})
            return
        }

        if count > int64(limit) {
            c.AbortWithStatusJSON(429, gin.H{"error": "Too Many Requests"})
            return
        }

        c.Next()
    }
}

以上代码演示了在15秒内只能有1次请求的限制，并将这个规则做成了gin框架的中间件。