-
Notifications
You must be signed in to change notification settings - Fork 101
/
ratelimit.go
114 lines (95 loc) · 2.42 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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package middleware
import (
"context"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/zhufuyi/sponge/pkg/gin/response"
rl "github.com/zhufuyi/sponge/pkg/shield/ratelimit"
)
// 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()})
}
}
// Timeout request time out
func Timeout(d time.Duration) gin.HandlerFunc {
if d < time.Millisecond {
return func(c *gin.Context) {}
}
return func(c *gin.Context) {
ctx, _ := context.WithTimeout(c.Request.Context(), d) //nolint
c.Request = c.Request.WithContext(ctx)
c.Next()
if ctx.Err() == context.DeadlineExceeded {
if c.Writer.Status() == 200 {
c.AbortWithStatus(http.StatusGatewayTimeout)
return
}
c.Abort()
}
}
}