42 lines
978 B
Go
42 lines
978 B
Go
|
package ratelimit
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"net/http"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/gin-gonic/gin"
|
|||
|
|
"go.uber.org/ratelimit"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
var rl = ratelimit.New(20, ratelimit.WithoutSlack)
|
|||
|
|
|
|||
|
|
func RateLimitMiddleware() gin.HandlerFunc {
|
|||
|
|
return func(c *gin.Context) {
|
|||
|
|
// 1. 给请求设置超时（如3秒，超过3秒未拿到令牌直接返回）
|
|||
|
|
ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
|
|||
|
|
defer cancel()
|
|||
|
|
|
|||
|
|
// 2. 用带超时的方式获取令牌（需自定义限流器的 Take 逻辑）
|
|||
|
|
ch := make(chan struct{}, 1)
|
|||
|
|
go func() {
|
|||
|
|
rl.Take() // 阻塞拿令牌
|
|||
|
|
ch <- struct{}{} // 拿到令牌后通知主协程
|
|||
|
|
}()
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case <-ch:
|
|||
|
|
// 3. 拿到令牌，继续执行
|
|||
|
|
c.Next()
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
// 4. 超时未拿到令牌，返回请求频繁
|
|||
|
|
c.JSON(http.StatusTooManyRequests, gin.H{
|
|||
|
|
"code": 429,
|
|||
|
|
"message": "请求过于频繁，请稍后重试",
|
|||
|
|
"reason": "等待令牌超时",
|
|||
|
|
})
|
|||
|
|
c.Abort()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|