【Golang】用官方rate包构造简单IP限流器

使用 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 middleware

import (
	"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] = limiter
	i.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() {
			// 触发限流，返回响应xxxx
			return
		}
		c.Next()
	}
}

4. 在 Gin 中使用中间件

最后，我们在 Gin 路由中使用这个中间件，并定期清理过期的 IP 限流器。

package main

import (
	"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 地址是否符合限流条件。