kpt-pasture/http/middleware/rate_limit.go

package middleware

import (
	"fmt"
	"time"

	"gitee.com/xuyiping_admin/pkg/apierr"
	"gitee.com/xuyiping_admin/pkg/xerr"
	"github.com/gin-gonic/gin"
	limit "github.com/yangxikun/gin-limit-by-key"
	"golang.org/x/time/rate"
)

// RateLimit limit rate by client ip
// limit 10 qps/client_ip and permit bursts of at most 10 tokens, and the limiter live time duration is 3 minutes
// handle exceed rate limit request
func RateLimit(opts ...*RateLimitOption) gin.HandlerFunc {
	var opt *RateLimitOption
	if len(opts) == 0 {
		opt = defaultRateLimitOption
	} else {
		opt = opts[0]
	}

	return limit.NewRateLimiter(func(c *gin.Context) string {
		return fmt.Sprintf("%s-%s", c.ClientIP(), authorization(c))
	}, func(c *gin.Context) (*rate.Limiter, time.Duration) {
		return rate.NewLimiter(rate.Every(opt.Interval), opt.Tokens), opt.LiveTime
	}, func(c *gin.Context) {
		apierr.AbortBadRequest(c, 429, xerr.Customf("request too many %s", c.ClientIP()))
	})
}

type RateLimitOption struct {
	Interval time.Duration
	Tokens   int
	LiveTime time.Duration
}

var defaultRateLimitOption = &RateLimitOption{
	Interval: 100 * time.Millisecond,
	Tokens:   10,
	LiveTime: 3 * time.Minute,
}

var testRateLimitOption = &RateLimitOption{
	Interval: 10 * time.Millisecond,
	Tokens:   1000,
	LiveTime: 1 * time.Minute,
}

func RateLimitOptionLow() *RateLimitOption {
	return &RateLimitOption{
		Interval: 1 * time.Second,
		Tokens:   10,
		LiveTime: 3 * time.Minute,
	}
}