You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
105 lines
2.4 KiB
105 lines
2.4 KiB
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
|
|
}
|
|
|