Go 实现接口请求次数限制

实现请求限制的原理主要基于对每个用户的请求进行计数和时间管理，从而在每个时间窗口内（例如每秒或每分钟）限制请求的数量。以下是实现这一原理的具体步骤和机制：

核心原理

1. 用户标识：使用唯一标识符（例如Token、用户ID）来区分不同的用户。这可以从HTTP请求头、Cookie或其他方式获取。
2. 计数器：为每个用户维护两个计数器，一个用于每分钟请求数（RPM），另一个用于每秒请求数（RPS）。
3. 时间窗口：使用时间戳记录上一次请求的时间，并根据当前时间与上一次请求时间的差值，决定是否重置计数器。
4. 限制检查：在每次请求到达时，检查当前计数器是否超过预设的限制（例如每秒最多3个请求，每分钟最多100个请求）。如果超过限制，返回错误响应；否则，更新计数器并允许请求通过。

具体实现步骤

1. 初始化用户请求信息：

   • 为每个用户创建一个结构体userRequestInfo，包括两个计数器和一个时间戳。
   • 使用全局映射表userRequests存储每个用户的请求信息。

2. 更新计数器：

   • 在每次请求到达时，根据当前时间与上一次请求时间的差值，更新每秒和每分钟的计数器。
   • 如果当前时间与上一次请求时间超过一秒，则重置每秒计数器；如果超过一分钟，则重置每分钟计数器。

3. 检查限制：

   • 检查更新后的计数器是否超过预设限制。
   • 如果超过限制，返回HTTP 429错误响应；否则，更新计数器并允许请求通过。

4. 中间件实现：

   • 使用中间件拦截所有请求，在实际处理请求前进行限流检查。
   • 如果通过限流检查，则继续处理请求；否则，返回错误响应。

代码示例

package main

import (
    "net/http"
    "sync"
    "time"

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

// 定义用户请求信息结构体
type userRequestInfo struct {
    mu              sync.Mutex
    requestsPerMin  int
    requestsPerSec  int
    lastRequestTime time.Time
}

// 用户请求信息映射表
var userRequests = make(map[string]*userRequestInfo)
var mu sync.Mutex

// 限制参数
const (
    maxRequestsPerMin = 100
    maxRequestsPerSec = 3
)

// 中间件函数
func rateLimitMiddleware() gin.HandlerFunc {
    return func(c *gin.Context) {
        token := c.GetHeader("Authorization")
        if token == "" {
            c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Token not provided"})
            return
        }

        mu.Lock()
        userInfo, exists := userRequests[token]
        if !exists {
            userInfo = &userRequestInfo{}
            userRequests[token] = userInfo
        }
        mu.Unlock()

        userInfo.mu.Lock()
        defer userInfo.mu.Unlock()

        now := time.Now()

        // 更新每秒请求数
        if now.Sub(userInfo.lastRequestTime) >= time.Second {
            userInfo.requestsPerSec = 0
        }

        // 更新每分钟请求数
        if now.Sub(userInfo.lastRequestTime) >= time.Minute {
            userInfo.requestsPerMin = 0
        }

        // 检查是否超过限制
        if userInfo.requestsPerSec >= maxRequestsPerSec {
            c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests per second"})
            return
        }

        if userInfo.requestsPerMin >= maxRequestsPerMin {
            c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests per minute"})
            return
        }

        // 更新请求数
        userInfo.requestsPerSec++
        userInfo.requestsPerMin++
        userInfo.lastRequestTime = now

        c.Next()
    }
}

func main() {
    r := gin.Default()

    // 使用中间件
    r.Use(rateLimitMiddleware())

    // 示例路由
    r.GET("/", func(c *gin.Context) {
        c.JSON(http.StatusOK, gin.H{"message": "Hello, World!"})
    })

    r.Run(":8080") // 启动服务器
}

详细解释

• 用户请求信息结构体：userRequestInfo包括了两个计数器和一个时间戳，用于记录每个用户的请求信息。
• 全局映射表：userRequests存储了所有用户的请求信息，使用Token作为键。
• 中间件函数：rateLimitMiddleware在每次请求到达时执行，检查并更新用户的请求计数器，并根据限制条件决定是否允许请求通过。
• 更新计数器和检查限制：在请求到达时，根据当前时间和上一次请求时间的差值，更新计数器，并检查是否超过限制。

通过这种方式，可以有效地限制每个用户的请求频率，防止服务器过载，提高系统的稳定性和安全性。