使用 Go 實現基于 IP 地址的限流機制

在高流量的服務中，限流是一個至關重要的環節。它不僅能防止服務因過多請求而崩潰，還能保證服務的公平性。在本篇文章中，我們將介紹如何利用 Go 語言中的 golang.org/x/time/rate 包實現一個基于 IP 地址的限流機制。

什么是 IP 限流？

IP 限流是指根據客戶端的 IP 地址，對每個 IP 地址發出的請求進行限制。這樣可以防止某個單獨的用戶或機器通過頻繁的請求耗盡服務器資源。

基于 rate.Limiter 實現 IP 限流

Go 語言的 golang.org/x/time/rate 包提供了一個非常實用的令牌桶限流算法，能夠輕松控制請求的速率。在本實現中，我們將基于客戶端 IP 地址，為每個 IP 地址創建一個獨立的限流器。

1. 設計思路

我們使用 map[string]*rate.Limiter 來存儲每個 IP 地址的限流器，并為每個 IP 地址設定一個令牌桶。我們還要為每個 IP 地址記錄最后一次請求的時間，并在一定時間（TTL）后清除過期的限流器。

2. 代碼實現

首先，我們創建一個 IPRateLimiter 結構體，該結構體維護了所有 IP 的限流器、過期時間以及并發安全的鎖。

package middlewareimport ("github.com/gin-gonic/gin""github.com/sirupsen/logrus""golang.org/x/time/rate""sync""time""git.nominee.com.cn/sectrainx/stx-common/pkg/response"
)// IPRateLimiter 定義了 IP 限流器
type IPRateLimiter struct {ips        map[string]*rate.Limiter // 存儲每個 IP 的限流器ipLastUsed map[string]time.Time     // 記錄每個 IP 地址最后一次使用的時間mu         *sync.RWMutex            // 確保并發安全r          rate.Limit               // 每秒生成令牌數b          int                      // 令牌桶容量ttl        time.Duration            // 限流器的過期時間
}// NewIPRateLimiter 創建一個新的 IP 限流器實例
func NewIPRateLimiter(r rate.Limit, b int, ttl time.Duration) *IPRateLimiter {i := &IPRateLimiter{ips:        make(map[string]*rate.Limiter),ipLastUsed: make(map[string]time.Time),mu:         &sync.RWMutex{},r:          r,b:          b,ttl:        ttl,}return i
}// AddIP 創建并添加一個新的 IP 限流器
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {i.mu.Lock()defer i.mu.Unlock()limiter := rate.NewLimiter(i.r, i.b)i.ips[ip] = limiteri.ipLastUsed[ip] = time.Now()	// 設置 IP 的最后使用時間return limiter
}// GetLimiter 獲取指定 IP 地址的限流器
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {i.mu.Lock()defer i.mu.Unlock()// 如果限流器存在，并且 IP 沒有過期，則返回現有的限流器if limiter, exists := i.ips[ip]; exists {if time.Since(i.ipLastUsed[ip]) < i.ttl {i.ipLastUsed[ip] = time.Now() // 更新最后訪問時間return limiter} else {// 如果限流器已經過期，刪除delete(i.ips, ip)delete(i.ipLastUsed, ip)}}// 不存在，創建新的限流器return i.AddIP(ip)
}// 定時清理過期的 IP 限流器
func (i *IPRateLimiter) CleanUpExpiredLimiters() {i.mu.Lock()defer i.mu.Unlock()for ip, lastUsed := range i.ipLastUsed {if time.Since(lastUsed) > i.ttl {delete(i.ips, ip)delete(i.ipLastUsed, ip)}}
}// 啟動一個定時器定期清理過期 IP 限流器
func (i *IPRateLimiter) StartCleanupRoutine() {go func() {for {time.Sleep(10 * time.Minute) // 每 10 分鐘執行一次清理i.CleanUpExpiredLimiters()}}()
}

3. 限流中間件

接下來，我們實現一個 RateLimitMiddleware 中間件，用于在每次請求時檢查客戶端的 IP 地址，并根據限流器的狀態判斷是否允許請求通過。

// IPRateLimitMiddleware 根據 IP 地址進行限流
func IPRateLimitMiddleware(limiter *IPRateLimiter) gin.HandlerFunc {return func(c *gin.Context) {ip := c.ClientIP()limiter := limiter.GetLimiter(ip)// 非阻塞方式檢查令牌if !limiter.Allow() {// 觸發限流，返回響應xxxxreturn}c.Next()}
}

4. 在 Gin 中使用中間件

最后，我們在 Gin 路由中使用這個中間件，并定期清理過期的 IP 限流器。

package mainimport ("github.com/gin-gonic/gin""log""time""your_project/middleware" // 引入你的中間件包
)func main() {// 創建一個 IP 限流器實例，設定每秒 10 個令牌，令牌桶容量為 100，TTL 為 1 小時ipLimiter := middleware.NewIPRateLimiter(10, 100, 1*time.Hour)// 啟動定時清理過期 IP 限流器ipLimiter.StartCleanupRoutine()// 創建 Gin 引擎r := gin.Default()// 將限流中間件應用到路由r.Use(middleware.RateLimitMiddleware2(ipLimiter))// 示例路由r.GET("/", func(c *gin.Context) {c.JSON(200, gin.H{"message": "Request successful!",})})// 啟動服務器if err := r.Run(":8080"); err != nil {log.Fatalf("Error starting server: %v", err)}
}

代碼解釋

1. IPRateLimiter 結構體：

   • ips 存儲每個 IP 地址的限流器。
   • ipLastUsed 記錄每個 IP 地址最后一次請求的時間。
   • r 和 b 分別是每秒生成令牌的速率和令牌桶的容量。
   • ttl 是過期時間，表示 IP 限流器的有效期。

2. NewIPRateLimiter： 初始化一個 IPRateLimiter 實例，設置速率、容量和過期時間。

3. GetLimiter： 獲取指定 IP 地址的限流器，如果限流器已經過期，會重新創建一個。

4. CleanUpExpiredLimiters： 定時清理過期的限流器。

5. RateLimitMiddleware2： 這是我們設計的限流中間件，檢查請求的 IP 地址是否符合限流條件。