/
ratelimit.go
92 lines (76 loc) · 1.95 KB
/
ratelimit.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
package middleware
import (
"net/http"
"time"
"github.com/zhufuyi/pkg/gin/response"
rl "github.com/zhufuyi/pkg/shield/ratelimit"
"github.com/gin-gonic/gin"
)
// ErrLimitExceed is returned when the rate limiter is
// triggered and the request is rejected due to limit exceeded.
var ErrLimitExceed = rl.ErrLimitExceed
// RateLimitOption set the rate limits rateLimitOptions.
type RateLimitOption func(*rateLimitOptions)
type rateLimitOptions struct {
window time.Duration
bucket int
cpuThreshold int64
cpuQuota float64
}
func defaultRatelimitOptions() *rateLimitOptions {
return &rateLimitOptions{
window: time.Second * 10,
bucket: 100,
cpuThreshold: 800,
}
}
func (o *rateLimitOptions) apply(opts ...RateLimitOption) {
for _, opt := range opts {
opt(o)
}
}
// WithWindow with window size.
func WithWindow(d time.Duration) RateLimitOption {
return func(o *rateLimitOptions) {
o.window = d
}
}
// WithBucket with bucket size.
func WithBucket(b int) RateLimitOption {
return func(o *rateLimitOptions) {
o.bucket = b
}
}
// WithCPUThreshold with cpu threshold
func WithCPUThreshold(threshold int64) RateLimitOption {
return func(o *rateLimitOptions) {
o.cpuThreshold = threshold
}
}
// WithCPUQuota with real cpu quota(if it can not collect from process correct);
func WithCPUQuota(quota float64) RateLimitOption {
return func(o *rateLimitOptions) {
o.cpuQuota = quota
}
}
// RateLimit an adaptive rate limiter middleware
func RateLimit(opts ...RateLimitOption) gin.HandlerFunc {
o := defaultRatelimitOptions()
o.apply(opts...)
limiter := rl.NewLimiter(
rl.WithWindow(o.window),
rl.WithBucket(o.bucket),
rl.WithCPUThreshold(o.cpuThreshold),
rl.WithCPUQuota(o.cpuQuota),
)
return func(c *gin.Context) {
done, err := limiter.Allow()
if err != nil {
response.Output(c, http.StatusTooManyRequests, err.Error())
c.Abort()
return
}
c.Next()
done(rl.DoneInfo{Err: c.Request.Context().Err()})
}
}