golang-cookbook

Appearance

qps ​

在使用 Gin 框架时，限制 QPS（Queries Per Second）可以通过中间件实现。以下是几种常见的方法：

方法 1：使用 time.Tick 和 channel 实现限流 ​

通过 time.Tick 和带缓冲的 channel 实现简单的令牌桶算法，限制每秒的请求数。

实现代码 ​

go
package main

import (
	"net/http"
	"time"

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

// 限流中间件
func rateLimiter(maxQPS int) gin.HandlerFunc {
	ticker := time.NewTicker(time.Second / time.Duration(maxQPS))
	defer ticker.Stop()

	return func(c *gin.Context) {
		select {
		case <-ticker.C:
			c.Next() // 允许请求通过
		default:
			c.JSON(http.StatusTooManyRequests, gin.H{
				"message": "Too many requests",
			})
			c.Abort() // 中断请求
		}
	}
}

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

	// 应用限流中间件，限制 QPS 为 10
	r.Use(rateLimiter(10))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	r.Run(":8080")
}

方法 2：使用第三方库 golang.org/x/time/rate ​

Go 标准库的扩展包 golang.org/x/time/rate 提供了更强大的限流功能，基于令牌桶算法。

实现代码 ​

go
package main

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

// 限流中间件
func rateLimiter(maxQPS int) gin.HandlerFunc {
	limiter := rate.NewLimiter(rate.Every(time.Second/time.Duration(maxQPS)), maxQPS)

	return func(c *gin.Context) {
		if !limiter.Allow() {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"message": "Too many requests",
			})
			c.Abort()
			return
		}
		c.Next()
	}
}

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

	// 应用限流中间件，限制 QPS 为 10
	r.Use(rateLimiter(10))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	r.Run(":8080")
}

方法 3：使用 Redis 实现分布式限流 ​

在分布式环境中，可以使用 Redis 实现全局限流。通过 Redis 的 INCR 和 EXPIRE 命令统计每秒的请求数。

实现代码 ​

go
package main

import (
	"context"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
)

var rdb *redis.Client

// 初始化 Redis 客户端
func initRedis() {
	rdb = redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
}

// 限流中间件
func rateLimiter(maxQPS int) gin.HandlerFunc {
	return func(c *gin.Context) {
		ctx := context.Background()
		key := "rate_limit:" + time.Now().Format("20060102150405") // 按秒分桶

		// 使用 INCR 增加计数
		count, err := rdb.Incr(ctx, key).Result()
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{
				"message": "Internal server error",
			})
			c.Abort()
			return
		}

		// 如果是第一次请求，设置过期时间
		if count == 1 {
			rdb.Expire(ctx, key, time.Second)
		}

		// 检查是否超过限制
		if count > int64(maxQPS) {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"message": "Too many requests",
			})
			c.Abort()
			return
		}

		c.Next()
	}
}

func main() {
	initRedis()

	r := gin.Default()

	// 应用限流中间件，限制 QPS 为 10
	r.Use(rateLimiter(10))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	r.Run(":8080")
}

方法 4：使用 github.com/ulule/limiter 库 ​

ulule/limiter 是一个功能强大的限流库，支持内存、Redis 等多种存储后端。

实现代码 ​

go
package main

import (
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/middleware/gin"
	"github.com/ulule/limiter/v3/drivers/store/memory"
)

func main() {
	// 定义限流规则：每秒 10 个请求
	rate := limiter.Rate{
		Period: 1 * time.Second,
		Limit:  10,
	}

	// 使用内存存储
	store := memory.NewStore()

	// 创建限流中间件
	middleware := ginlimiter.NewMiddleware(limiter.New(store, rate))

	r := gin.Default()

	// 应用限流中间件
	r.Use(middleware)

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	r.Run(":8080")
}

总结 ​

  • 单机限流：推荐使用 golang.org/x/time/rate 或 time.Tick 实现。
  • 分布式限流：推荐使用 Redis 或 ulule/limiter 库。
  • 性能优化：在高并发场景下，避免频繁的内存分配和锁竞争，尽量使用高效的算法和存储。

根据实际需求选择合适的限流方案，并结合监控和日志系统，确保服务的稳定性和可扩展性。