logo

使用Redis实现Golang API限流

Published on

限流简单来说就是一种限制用户或客户端在指定时间范围内能够向API发送请求次数的技术。你可能曾经在访问天气API或笑话API时遇到过"超出请求限制"的提示信息。关于为什么要对API进行限流，有很多观点，但一些重要的原因包括：确保公平使用、增强安全性、防止资源过载等。

在本文中，我们将使用Gin框架创建一个Golang HTTP服务器，并使用Redis对某个端点实现限流功能，记录指定时间范围内来自同一IP的请求总数。当请求超过我们设定的限制时，将返回错误消息。

如果你对Gin和Redis还不了解：Gin是一个用Golang编写的Web框架，它可以帮助我们用较少的代码创建简单快速的服务器。Redis是一个内存键值数据存储系统，可以用作数据库或缓存。

前提条件

  • 熟悉Golang、Gin和Redis
  • 一个Redis实例（可以使用Docker或远程机器）

开始使用

首先通过运行 go mod init <github路径> 来初始化项目，例如：go mod init github.com/Pradumnasaraf/go-redis。

让我们先用Gin框架创建一个简单的HTTP服务器，然后再实现限流逻辑。你可以复制下面的代码，这是一个很基础的示例。当我们访问 /message 端点时，服务器会返回一条消息。

复制下面的代码后，运行 go mod tidy 来自动安装我们导入的包。

package main

import (
    "github.com/gin-gonic/gin"
)

func main() {
    r := gin.Default()
    r.GET("/message", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "You can make more requests",
        })
    })
    r.Run(":8081") //监听并服务于localhost:8081
}

我们可以通过在终端执行 go run main.go 来运行服务器

GC Performance

测试时，我们可以访问 localhost:8081/message，就能在浏览器中看到消息。

GC Performance

现在我们的服务器已经运行起来了，让我们为 /message 路由设置限流功能。我们将使用 go-redis/redis_rate 包。感谢这个包的作者，我们不需要从头编写处理和检查限制的逻辑。它会为我们处理所有繁重的工作。

以下是实现限流功能后的完整代码。我们会逐步理解其中的每个部分：

package main

import (
    "context"
    "errors"
    "net/http"
    "github.com/gin-gonic/gin"
    "github.com/go-redis/redis_rate/v10"
    "github.com/redis/go-redis/v9"
)

var (
    rdb     *redis.Client
    limiter *redis_rate.Limiter
)

func initRedis() {
    rdb = redis.NewClient(&redis.Options{
        Addr: "localhost:6379",
    })
    limiter = redis_rate.NewLimiter(rdb)
}

func main() {
    // 初始化Redis客户端和限流器
    initRedis()
    defer rdb.Close()

    r := gin.Default()
    r.GET("/message", func(c *gin.Context) {
        err := rateLimiter(c.ClientIP())
        if err != nil {
            c.JSON(http.StatusTooManyRequests, gin.H{
                "message": "you have hit the limit",
            })
            return
        }
        c.JSON(http.StatusOK, gin.H{
            "message": "You can make more requests",
        })
    })
    r.Run(":8081")
}

func rateLimiter(clientIP string) error {
    ctx := context.Background()

    res, err := limiter.Allow(ctx, clientIP, redis_rate.PerMinute(10))
    if err != nil {
        return err
    }
    if res.Remaining == 0 {
        return errors.New("Rate limit exceeded")
    }

    return nil
}

让我们首先看 initRedis() 函数。这个函数会在应用程序启动时创建一个Redis客户端和限流器实例。这样，我们就不需要每次都创建新的实例。我们创建了全局变量 rdb 来存储redis实例，limiter 来存储限流器实例。

现在来理解 rateLimiter() 函数。这个函数需要一个参数，即请求的IP地址，我们可以在main函数中通过 c.ClientIP() 获取。如果达到限制就返回错误，否则返回nil。这里的大部分代码都是从官方GitHub仓库获取的样板代码。需要重点关注的是 limiter.Allow() 函数，Addr：获取 Redis 实例的 URL 路径值。我使用 Docker 在本地运行它。您可以使用任何内容，请确保相应地替换 URL。：

res, err := limiter.Allow(ctx, clientIP, redis_rate.PerMinute(10))

它接受三个参数：

  • 第一个是上下文ctx
  • 第二个是Redis数据库的键（Key）
  • 第三个是限制值

该函数将clientIP地址作为键，默认限制作为值存储，并在收到请求时递减。之所以使用这种结构，是因为Redis数据库需要唯一标识和唯一键来存储键值对数据，而每个IP地址本身就是唯一的，这就是为什么我们使用IP地址而不是用户名等。

第三个参数 redis_rate.PerMinute(10) 可以根据需要修改，我们可以设置 PerSecond、PerHour 等限制，括号内的值表示每分钟/秒/小时可以发出多少请求。在我们的例子中是每分钟10次。是的，设置起来就是这么简单。

最后，我们通过 res.Remaining 检查是否还有剩余配额。如果为零，我们将返回一个错误消息，否则返回nil。例如，你也可以使用 res.Limit.Rate 检查限制率等。你可以进一步探索这些功能。需要注意的是，这只是一个如何将这两个部分组合在一起的示例，因为我们只有一个路由，所以没有使用中间件，但如果我们有10个或100个路由呢？

再来看main()函数中的实现部分：

func main() {
    // 初始化Redis客户端和限流器
    initRedis()
    defer rdb.Close()

    r := gin.Default()
    r.GET("/message", func(c *gin.Context) {
        err := rateLimiter(c.ClientIP())
        if err != nil {
            c.JSON(http.StatusTooManyRequests, gin.H{
                "message": "you have hit the limit",
            })
            return
        }
        c.JSON(http.StatusOK, gin.H{
            "message": "You can make more requests",
        })
    })
    r.Run(":8081")
}

main()函数中的大部分内容都保持不变。我们调用 initRedis() 函数来初始化Redis客户端和限流器，然后使用defer在应用程序退出时关闭redis客户端。在 /message 路由中，每次访问该路由时，我们都会调用 rateLimit() 函数并传入ClientIP地址，并将返回值（错误）存储在err变量中。如果有错误，我们将返回429状态码（即 http.StatusTooManyRequests）和消息 "message": "You have hit the limit"。如果用户还有剩余限制且 rateLimit() 没有返回错误，它将像之前一样正常工作并处理请求。

以上就是所有的解释。让我们来测试一下。重新运行服务器执行相同的命令。第一次我们会看到和之前相同的消息。现在刷新浏览器10次（因为我们设置了每分钟10次的限制），你就会在浏览器中看到错误消息。

GC Performance

我们也可以通过查看终端中的日志来验证这一点。Gin默认提供了很好的日志功能。一分钟后，我们的限制配额将会恢复。

GC Performance