/
limiter.go
105 lines (87 loc) · 2.36 KB
/
limiter.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
package polaris
import (
"time"
"github.com/go-kratos/aegis/ratelimit"
"github.com/polarismesh/polaris-go"
"github.com/polarismesh/polaris-go/pkg/model"
)
type (
// LimiterOption function for polaris limiter
LimiterOption func(*limiterOptions)
)
type limiterOptions struct {
// required, polaris limit namespace
namespace string
// required, polaris limit service name
service string
// optional, polaris limit request timeout
// max value is (1+RetryCount) * Timeout
timeout time.Duration
// optional, polaris limit retryCount
// init by polaris config
retryCount int
// optional, request limit quota
token uint32
}
// WithLimiterNamespace with limiter namespace.
func WithLimiterNamespace(namespace string) LimiterOption {
return func(o *limiterOptions) {
o.namespace = namespace
}
}
// WithLimiterService with limiter service.
func WithLimiterService(service string) LimiterOption {
return func(o *limiterOptions) {
o.service = service
}
}
// WithLimiterTimeout with limiter arguments.
func WithLimiterTimeout(timeout time.Duration) LimiterOption {
return func(o *limiterOptions) {
o.timeout = timeout
}
}
// WithLimiterRetryCount with limiter retryCount.
func WithLimiterRetryCount(retryCount int) LimiterOption {
return func(o *limiterOptions) {
o.retryCount = retryCount
}
}
// WithLimiterToken with limiter token.
func WithLimiterToken(token uint32) LimiterOption {
return func(o *limiterOptions) {
o.token = token
}
}
type Limiter struct {
// polaris limit api
limitAPI polaris.LimitAPI
opts limiterOptions
}
// init quotaRequest
func buildRequest(opts limiterOptions) polaris.QuotaRequest {
quotaRequest := polaris.NewQuotaRequest()
quotaRequest.SetNamespace(opts.namespace)
quotaRequest.SetRetryCount(opts.retryCount)
quotaRequest.SetService(opts.service)
quotaRequest.SetTimeout(opts.timeout)
quotaRequest.SetToken(opts.token)
return quotaRequest
}
// Allow interface impl
func (l *Limiter) Allow(method string, argument ...model.Argument) (ratelimit.DoneFunc, error) {
request := buildRequest(l.opts)
request.SetMethod(method)
for _, arg := range argument {
request.AddArgument(arg)
}
resp, err := l.limitAPI.GetQuota(request)
if err != nil {
// ignore err
return func(ratelimit.DoneInfo) {}, nil
}
if resp.Get().Code == model.QuotaResultOk {
return func(ratelimit.DoneInfo) {}, nil
}
return nil, ratelimit.ErrLimitExceed
}