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()
|
||
}
|
||
}
|
||
}
|